add test to make sure we are properly checking for max_num_elements

FAIR-Chem · Aug 20, 2024 · a27e658 · a27e658
1 parent f47e020
commit a27e658
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/tests/core/e2e/test_s2ef.py b/tests/core/e2e/test_s2ef.py
@@ -170,6 +170,26 @@ def test_use_pbc_single(self, configs, tutorial_val_src, torch_deterministic):
                 input_yaml=configs["equiformer_v2"],
             )
 
+    def test_max_num_atoms(self, configs, tutorial_val_src, torch_deterministic):
+        with tempfile.TemporaryDirectory() as tempdirname:
+            tempdir = Path(tempdirname)
+            extra_args = {"seed": 0}
+            with pytest.raises(AssertionError):
+                _ = _run_main(
+                    rundir=str(tempdir),
+                    update_dict_with={
+                        "optim": {"max_epochs": 1},
+                        "model": {"max_num_elements": 2},
+                        "dataset": oc20_lmdb_train_and_val_from_paths(
+                            train_src=str(tutorial_val_src),
+                            val_src=str(tutorial_val_src),
+                            test_src=str(tutorial_val_src),
+                        ),
+                    },
+                    update_run_args_with=extra_args,
+                    input_yaml=configs["equiformer_v2"],
+                )
+
     @pytest.mark.parametrize(
         ("world_size", "ddp"),
         [