safety

mace_jax/run_train.py

@@ -32,7 +32,10 @@ def main():
     train_loader, valid_loader, test_loader, atomic_energies_dict, r_max = datasets()
 
     model_fn, params, num_message_passing = model(
-        r_max, atomic_energies_dict, train_loader.graphs, initialize_seed=seed
+        r_max=r_max,
+        atomic_energies_dict=atomic_energies_dict,
+        train_graphs=train_loader.graphs,
+        initialize_seed=seed,
     )
 
     params = reload(params)

mace_jax/tools/gin_model.py

@@ -80,11 +80,11 @@ def __call__(self, node_specie: jnp.ndarray) -> e3nn.IrrepsArray:
 
 @gin.configurable
 def model(
+    *,
     r_max: float,
     atomic_energies_dict: Dict[int, float] = None,
     train_graphs: List[jraph.GraphsTuple] = None,
     initialize_seed: Optional[int] = None,
-    *,
     scaling: Callable = None,
     atomic_energies: Union[str, np.ndarray, Dict[int, float]] = None,
     avg_num_neighbors: float = "average",
@@ -166,11 +166,12 @@ def model(
 
     # check that num_species is consistent with the dataset
     if z_table is None:
-        for graph in train_graphs:
-            if not np.all(graph.nodes.species < num_species):
-                raise ValueError(
-                    f"max(graph.nodes.species)={np.max(graph.nodes.species)} >= num_species={num_species}"
-                )
+        if train_graphs is not None:
+            for graph in train_graphs:
+                if not np.all(graph.nodes.species < num_species):
+                    raise ValueError(
+                        f"max(graph.nodes.species)={np.max(graph.nodes.species)} >= num_species={num_species}"
+                    )
     else:
         if max(z_table.zs) >= num_species:
             raise ValueError(