From 83a8ae75b2e72b8b36758318a1f5c87a3d68589a Mon Sep 17 00:00:00 2001 From: kenko911 Date: Sat, 28 Sep 2024 15:26:25 -0700 Subject: [PATCH] fix pytest --- m3gnet/models/_dynamics.py | 2 +- m3gnet/models/tests/test_model.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/m3gnet/models/_dynamics.py b/m3gnet/models/_dynamics.py index 6274518..86d158b 100644 --- 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 index 4b5f898..228c48a 100644 --- 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])