Merge pull request #276 from lory-w/275-ase-missing-data

add missing data in format atoms
IntelLabs · Aug 19, 2024 · 262a094 · 262a094
2 parents 46bb7db + d0c23ad
commit 262a094
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/matsciml/interfaces/ase/base.py b/matsciml/interfaces/ase/base.py
@@ -212,6 +212,9 @@ def _format_atoms(self, atoms: Atoms) -> DataDict:
         data_dict["pos"] = pos
         data_dict["atomic_numbers"] = atomic_numbers
         data_dict["cell"] = cell
+        # ptr and batch are usually expected by MACE even if it's a single graph
+        data_dict["ptr"] = torch.tensor([0])
+        data_dict["batch"] = torch.zeros((pos.size(0)))
         return data_dict
 
     def _format_pipeline(self, atoms: Atoms) -> DataDict: