pass is_molecule_mod to AtomInput

Author: lucidrains

alphafold3_pytorch/inputs.py

@@ -1075,24 +1075,25 @@ def molecule_lengthed_molecule_input_to_atom_input(mol_input: MoleculeLengthMole
     # atom input
 
     atom_input = AtomInput(
-        atom_inputs=atom_inputs_tensor,
-        atompair_inputs=atompair_inputs,
-        molecule_atom_lens=atoms_per_token,
-        molecule_ids=molecule_ids,
-        molecule_atom_indices=molecule_atom_indices,
-        distogram_atom_indices=distogram_atom_indices,
-        missing_atom_mask=missing_atom_mask,
-        additional_token_feats=additional_token_feats,
-        additional_molecule_feats=additional_molecule_feats,
-        is_molecule_types=is_molecule_types,
+        atom_inputs = atom_inputs_tensor,
+        atompair_inputs = atompair_inputs,
+        molecule_atom_lens = atoms_per_token,
+        molecule_ids = molecule_ids,
+        molecule_atom_indices = molecule_atom_indices,
+        distogram_atom_indices = distogram_atom_indices,
+        missing_atom_mask = missing_atom_mask,
+        additional_token_feats = additional_token_feats,
+        additional_molecule_feats = additional_molecule_feats,
+        is_molecule_mod = is_molecule_mod,
+        is_molecule_types = is_molecule_types,
         msa = msa,
         templates = templates,
-        atom_pos=atom_pos,
-        token_bonds=token_bonds,
-        atom_parent_ids=i.atom_parent_ids,
-        atom_ids=atom_ids,
-        atompair_ids=atompair_ids,
-        chains=chains,
+        atom_pos = atom_pos,
+        token_bonds = token_bonds,
+        atom_parent_ids = i.atom_parent_ids,
+        atom_ids = atom_ids,
+        atompair_ids = atompair_ids,
+        chains = chains,
     )
 
     return atom_input

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.2.88"
+version = "0.2.89"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }