add head_key=head_name to info dict of atoms regardless of isolated_a…

…tom_config_value
ACEsuit · Nov 9, 2024 · 0e5d72b · 0e5d72b
1 parent c1184eb
commit 0e5d72b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/mace/data/utils.py b/mace/data/utils.py
@@ -265,7 +265,6 @@ def load_from_xyz(
         atoms_without_iso_atoms = []
 
         for idx, atoms in enumerate(atoms_list):
-            atoms.info[head_key] = head_name
             isolated_atom_config = (
                 len(atoms) == 1 and atoms.info.get("config_type") == "IsolatedAtom"
             )
@@ -288,6 +287,9 @@ def load_from_xyz(
         if not keep_isolated_atoms:
             atoms_list = atoms_without_iso_atoms
 
+    for atoms in atoms_list:
+        atoms.info[head_key] = head_name
+
     configs = config_from_atoms_list(
         atoms_list,
         config_type_weights=config_type_weights,