[Bug]: Retrained M3Gnet potential cannot be used

Email (Optional)

No response

Version

0.9.1

Which OS(es) are you using?

[ ] MacOS
[ ] Windows
[X] Linux

What happened?

I retrained M3GNet-MP-2021.2.8-PES potentials following tutorial from https://matgl.ai/tutorials%2FTraining%20a%20M3GNet%20Potential%20with%20PyTorch%20Lightning.html using collected structures, energy, forces and stresses data from our structure relaxations. When I attempt to relax any structure using that potential, I got error AttributeError: 'M3GNet' object has no attribute 'calc_stresses'

Code snippet

from __future__ import annotations
.
import warnings

from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from pymatgen.core import Lattice, Structure
from pymatgen.io.ase import AseAtomsAdaptor

import matgl
from matgl.ext.ase import M3GNetCalculator, MolecularDynamics, Relaxer

warnings.simplefilter("ignore")

pot = matgl.load_model("MyRetrainedPotential") # this was put into ~/.cache/matgl after training
relaxer = Relaxer(potential=pot) # this produces following error:

Log output

Cell In[14], line 1
----> 1 relaxer = Relaxer(potential=pot)

File ~/miniconda3/envs/mgl/lib/python3.9/site-packages/matgl/ext/ase.py:211, in Relaxer.__init__(self, potential, state_attr, optimizer, relax_cell, stress_weight)
    200 """
    201 Args:
    202     potential (Potential): a M3GNet potential, a str path to a saved model or a short name for saved model
   (...)
    208     stress_weight (float): conversion factor from GPa to eV/A^3.
    209 """
    210 self.optimizer: Optimizer = OPTIMIZERS[optimizer.lower()].value if isinstance(optimizer, str) else optimizer
--> 211 self.calculator = M3GNetCalculator(
    212     potential=potential,
    213     state_attr=state_attr,
    214     stress_weight=stress_weight,  # type: ignore
    215 )
    216 self.relax_cell = relax_cell
    217 self.potential = potential

File ~/miniconda3/envs/mgl/lib/python3.9/site-packages/matgl/ext/ase.py:146, in M3GNetCalculator.__init__(self, potential, state_attr, stress_weight, **kwargs)
    144 super().__init__(**kwargs)
    145 self.potential = potential
--> 146 self.compute_stress = potential.calc_stresses
    147 self.compute_hessian = potential.calc_hessian
    148 self.stress_weight = stress_weight

File ~/miniconda3/envs/mgl/lib/python3.9/site-packages/torch/nn/modules/module.py:1695, in Module.__getattr__(self, name)
   1693     if name in modules:
   1694         return modules[name]
-> 1695 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

AttributeError: 'M3GNet' object has no attribute 'calc_stresses'

Code of Conduct

[X] I agree to follow this project's Code of Conduct

materialsvirtuallab / matgl

[Bug]: Retrained M3Gnet potential cannot be used #208