fix nequip-train edge case with no train run type

mir-group · Jan 3, 2025 · c8d8fc5 · c8d8fc5
1 parent 1c2295c
commit c8d8fc5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nequip/scripts/train.py b/nequip/scripts/train.py
@@ -53,8 +53,8 @@ def main(config: DictConfig) -> None:
 
     # ensure only single train at most, to protect restart and checkpointing logic later
     assert (
-        sum([run_type == "train" for run_type in runs]) == 1
-    ), "only a single `train` instance can be present in `run`"
+        sum([run_type == "train" for run_type in runs]) <= 1
+    ), "only up to a single `train` instance can be present in `run`"
 
     # ensure that the relevant metrics are present
     if "train" in runs: