Update inputs.py (#196)

Author: amorehead

alphafold3_pytorch/inputs.py

@@ -3160,9 +3160,7 @@ def pdb_input_to_molecule_input(
 
     # handle `atom_indices_for_frame` for the PAE
 
-    atom_indices_for_frame = tensor(
-        [default(indices, (-1, -1, -1)) for indices in atom_indices_for_frame]
-    )
+    atom_indices_for_frame = tensor(atom_indices_for_frame)
 
     # build offsets for all indices
 
@@ -3199,6 +3197,12 @@ def pdb_input_to_molecule_input(
     # craft ligand frame offsets
     atom_indices_for_ligand_frame = torch.zeros_like(atom_indices_for_frame)
     for ligand_frame_index in torch.where(is_ligand_frame)[0]:
+        if (atom_indices_for_frame[ligand_frame_index] == -1).any():
+            atom_indices_for_ligand_frame[ligand_frame_index] = atom_indices_for_frame[
+                ligand_frame_index
+            ]
+            continue
+
         global_atom_indices = torch.gather(
             atom_indices_offsets, 0, atom_indices_for_frame[ligand_frame_index]
         )