There must be a better way to implement this function, but I couldn’t find any similar solution. Can somebody help with this? x is of shape (500,6,4).
def myFcn(x):
cos_x2 = torch.cos(x[:,2])
sin_x2 = torch.sin(x[:,2])
cos_x5 = torch.cos(x[:,5])
sin_x5 = torch.sin(x[:,5])
out = torch.ones([x.shape[0], 6, 4])

for i in range(x.shape[0]):
  out[i,:,:] = torch.Tensor([
    [cos_x2[i], 0.0, 0.0, 0.0],
    [sin_x2[i], 0.0, 0.0, 0.0],
    [0.0, 1.0, 0.0, 0.0],
    [0.0, 0.0, cos_x5[i], 0.0],
    [0.0, 0.0, sin_x5[i], 0.0],
    [0.0, 0.0, 0.0, 1.0]])

return out