fix pytest

materialsvirtuallab · Sep 28, 2024 · 83a8ae7 · 83a8ae7
1 parent 0409a70
commit 83a8ae7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/m3gnet/models/_dynamics.py b/m3gnet/models/_dynamics.py
@@ -25,9 +25,9 @@
 from pymatgen.core.structure import Molecule, Structure
 from pymatgen.io.ase import AseAtomsAdaptor
 
+from m3gnet.layers import AtomRef
 from ._base import Potential
 from ._m3gnet import M3GNet
-from m3gnet.layers import AtomRef
 
 OPTIMIZERS = {
     "FIRE": FIRE,

diff --git a/m3gnet/models/tests/test_model.py b/m3gnet/models/tests/test_model.py
@@ -20,7 +20,7 @@ def setUpClass(cls) -> None:
         cls.atoms = Atoms(["Mo", "Mo"], [[0, 0, 0], [0.5, 0.5, 0.5]], cell=np.eye(3) * 3.30, pbc=True)
         cls.single_atoms = Structure(Lattice.cubic(6.0), ["Mo"], [[0, 0, 0]])
 
-        cls.mfi_model = M3GNet(is_intensive=False, state_embedding_dim=16, n_type_states=2)
+        cls.mfi_model = M3GNet(is_intensive=False, state_embedding_dim=16, n_state_types=2)
         cls.mfi_potential = Potential(model=cls.mfi_model)
         cls.state_attr = np.array([1])