fix pylint

materialsvirtuallab · Sep 28, 2024 · a949740 · a949740
1 parent 3113f77
commit a949740
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/m3gnet/models/_dynamics.py b/m3gnet/models/_dynamics.py
@@ -27,6 +27,7 @@
 
 from ._base import Potential
 from ._m3gnet import M3GNet
+from ._m3gnet.layers import AtomRef
 
 OPTIMIZERS = {
     "FIRE": FIRE,
@@ -96,7 +97,7 @@ def calculate(
         graph = self.potential.graph_converter(atoms, self.state_attr)
         graph_list = graph.as_tf().as_list()
         results = self.potential.get_efs_tensor(graph_list, include_stresses=self.compute_stress)
-
+        offset = None
         if self.element_refs is not None:
             offset = self.element_refs(graph_list)