add .long for group_idx

dauparas · May 7, 2024 · 091ab1f · 091ab1f
1 parent 84614f9
commit 091ab1f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sc_utils.py b/sc_utils.py
@@ -208,12 +208,12 @@ def make_torsion_features(feature_dict, repack_everything=True):
         S_af2,
         torch.tensor(restype_rigid_group_default_frame, device=device),
     )
-
+    
     xyz14_noised = feats.frames_and_literature_positions_to_atom14_pos(
         pred_frames,
         S_af2,
         torch.tensor(restype_rigid_group_default_frame, device=device),
-        torch.tensor(restype_atom14_to_rigid_group, device=device),
+        torch.tensor(restype_atom14_to_rigid_group, device=device).long(),
         torch.tensor(restype_atom14_mask, device=device),
         torch.tensor(restype_atom14_rigid_group_positions, device=device),
     )