fix post merge

FAIR-Chem · Aug 21, 2024 · da81a89 · da81a89
1 parent 2878bf6
commit da81a89
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/src/fairchem/core/models/equiformer_v2/equiformer_v2.py b/src/fairchem/core/models/equiformer_v2/equiformer_v2.py
@@ -397,6 +397,9 @@ def forward(self, data: Batch) -> dict[str, torch.Tensor]:
         self.dtype = data.pos.dtype
         self.device = data.pos.device
         atomic_numbers = data.atomic_numbers.long()
+        assert (
+            atomic_numbers.max().item() < self.max_num_elements
+        ), "Atomic number exceeds that given in model config"
         graph = self.generate_graph(
             data,
             enforce_max_neighbors_strictly=self.enforce_max_neighbors_strictly,

diff --git a/tests/core/e2e/test_s2ef.py b/tests/core/e2e/test_s2ef.py
@@ -179,15 +179,15 @@ def test_max_num_atoms(self, configs, tutorial_val_src, torch_deterministic):
                     rundir=str(tempdir),
                     update_dict_with={
                         "optim": {"max_epochs": 1},
-                        "model": {"max_num_elements": 2},
+                        "model": {"backbone": {"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"],
+                    input_yaml=configs["equiformer_v2_hydra"],
                 )
 
     @pytest.mark.parametrize(