Model restored from checkpoint predicts wrong energies

[dembart]

Dear m3gnet developers,

I'm getting the wrong energies when using a potential recovered from the callbacks folder. The MAE energy after training is ~ 0.005 eV/atom, but after the restoration of the model, the potential predicts completely different energies (~ 10 times less). However the forces are ok. Do I do something incorrectly while loading weights (I am not so familiar with tensorflow)? Could you suggest any solution? The exmaple of loading weights is below:

m3gnet = M3GNet(is_intensive=False)
folder= 'callbacks/'
latest = tf.train.latest_checkpoint(os.path.dirname(folder))
m3gnet.load_weights(latest)
potential = Potential(model=m3gnet)