compute rot 6d does not depend on cuda

src/StructDiffusion/diffusion/pose_conversion.py

@@ -45,11 +45,6 @@ def get_diffusion_variables_from_H(poses):
 
 def get_struct_objs_poses(x):
 
-    on_gpu = x.is_cuda
-    if not on_gpu:
-        x = x.cuda()
-
-    # assert x.is_cuda, "compute_rotation_matrix_from_ortho6d requires input to be on gpu"
     device = x.device
 
     # important: the noisy x can go out of bounds
@@ -72,10 +67,6 @@ def get_struct_objs_poses(x):
     struct_pose = x_full[:, 0].unsqueeze(1) # B, 1, 4, 4
     pc_poses_in_struct = x_full[:, 1:] # B, N, 4, 4
 
-    if not on_gpu:
-        struct_pose = struct_pose.cpu()
-        pc_poses_in_struct = pc_poses_in_struct.cpu()
-
     return struct_pose, pc_poses_in_struct
 
 

src/StructDiffusion/utils/rotation_continuity.py

@@ -21,7 +21,7 @@ def compute_pose_from_rotation_matrix(T_pose, r_matrix):
 def normalize_vector( v, return_mag =False):
     batch=v.shape[0]
     v_mag = torch.sqrt(v.pow(2).sum(1))# batch
-    v_mag = torch.max(v_mag, torch.autograd.Variable(torch.FloatTensor([1e-8]).cuda()))
+    v_mag = torch.max(v_mag, torch.autograd.Variable(torch.FloatTensor([1e-8]).to(v.device)))
     v_mag = v_mag.view(batch,1).expand(batch,v.shape[1])
     v = v/v_mag
     if(return_mag==True):