Skip to content

Commit

Permalink
put equiformer model in eval mode when testing
Browse files Browse the repository at this point in the history
  • Loading branch information
curtischong committed Aug 27, 2024
1 parent 5af7283 commit 7c237d6
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions tests/core/models/test_equiformer_v2_deprecated.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
from fairchem.core.common.utils import load_state_dict, setup_imports
from fairchem.core.datasets import data_list_collater
from fairchem.core.preprocessing import AtomsToGraphs
import logging


@pytest.fixture(scope="class")
Expand Down Expand Up @@ -102,6 +103,7 @@ def _load_model():
# so we explicitly set the number of layers to 1 (instead of all 8).
# The other alternative is to have different snapshots for mac vs. linux.
model.num_layers = 1
model.eval()
return model


Expand Down

0 comments on commit 7c237d6

Please sign in to comment.