Hi Ted,
I tried your code, it works. But how can I run it on GPU?
I added below

device_gpu = torch.device(“cuda”) #to gpu
data = torch.rand((2, 64, 256, 256), device=device_gpu).float()

but  I got error when runing rotate funciton, RuntimeError
Traceback (most recent call last)
 in 
----> 1 rotated_data = rotate(data, rot_mat)
2 print(data.shape)
3 print(rotated_data.shape)
4 print(torch.mean(rotated_data - data))
 in rotate(input_tensor, rotation_matrix)
14       locations_3d = get_3d_locations(d, h, w, device_)
15       # rotate target positions to the source coordinate
—> 16       rotated_3d_positions = torch.bmm(rotation_matrix.view(1, 3, 3).expand(dhw, 3, 3), locations_3d).view(1, d,h,w,3)
17       rot_locs = torch.split(rotated_3d_positions, split_size_or_sections=1, dim=4)
18       # change the range of x,y,z locations to [-1,1]
RuntimeError: Expected tensor to have CPU Backend, but got tensor with CUDA Backend (while checking arguments for bmm)