fix bug test set stress key

ACEsuit · May 17, 2024 · 4a4b12b · 4a4b12b
1 parent e6a6869
commit 4a4b12b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/mace/data/utils.py b/mace/data/utils.py
@@ -215,7 +215,7 @@ def load_from_xyz(
                 atoms.info["REF_energy"] = None
     if forces_key == "forces":
         logging.info(
-            "Since ASE version 3.23.0b1, using energy_key 'forces' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_forces'. You need to use --forces_key='REF_forces', to tell the key name chosen."
+            "Since ASE version 3.23.0b1, using force_key 'forces' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_forces'. You need to use --forces_key='REF_forces', to tell the key name chosen."
         )
         forces_key = "REF_forces"
         for atoms in atoms_list:
@@ -226,7 +226,7 @@ def load_from_xyz(
                 atoms.arrays["REF_forces"] = None
     if stress_key == "stress":
         logging.info(
-            "Since ASE version 3.23.0b1, using energy_key 'stress' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_stress'. You need to use --stress_key='REF_stress', to tell the key name chosen."
+            "Since ASE version 3.23.0b1, using stress_key 'stress' is no longer safe when communicating between MACE and ASE. We recommend using a different key, rewriting energies to 'REF_stress'. You need to use --stress_key='REF_stress', to tell the key name chosen."
         )
         stress_key = "REF_stress"
         for atoms in atoms_list:

diff --git a/mace/tools/scripts_utils.py b/mace/tools/scripts_utils.py
@@ -91,6 +91,8 @@ def get_dataset_from_xyz(
             energy_key=energy_key,
             forces_key=forces_key,
             dipole_key=dipole_key,
+            stress_key=stress_key,
+            virials_key=virials_key,
             charges_key=charges_key,
             extract_atomic_energies=False,
         )