Merge branch 'main' into step

ACEsuit · Feb 2, 2023 · f994f3e · f994f3e
2 parents 3ded1be + f5bbc34
commit f994f3e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/configs/aspirin_small.gin b/configs/aspirin_small.gin
@@ -8,7 +8,7 @@ logs.directory = "results"
 
 datasets.train_path = "data/rmd17_aspirin_train.xyz"  # in units of eV and Å
 datasets.test_path = "data/rmd17_aspirin_test.xyz"
-datasets.num_train = 150
+datasets.train_num = 150
 datasets.r_max = 5.0
 datasets.valid_fraction = 0.208
 datasets.n_node = 512

diff --git a/mace_jax/tools/gin_datasets.py b/mace_jax/tools/gin_datasets.py
@@ -14,7 +14,7 @@ def datasets(
     r_max: float,
     train_path: str,
     config_type_weights: Dict = None,
-    num_train: int = None,
+    train_num: int = None,
     valid_path: str = None,
     valid_fraction: float = None,
     valid_num: int = None,
@@ -46,7 +46,7 @@ def datasets(
         energy_key=energy_key,
         forces_key=forces_key,
         extract_atomic_energies=True,
-        num_configs=num_train,
+        num_configs=train_num,
         prefactor_stress=prefactor_stress,
         remap_stress=remap_stress,
     )