Question about loading the model after fine-tunning

[M9JS]

Hello guys, could someone kindly tell me how to load the model after fine-tunning?
I have a data-set consists of 90 structure snapshots sampled from a MD simulation, and calculated with vasp. As I follow the instruction in the notebook "fine_tuning.ipynb", I successfully labeled all of them and feed them to the model for fine-tunning. I set a lower learning_rate and larger epochs for better tunning.

# Optionally fix the weights of some layers
for layer in [
    chgnet.atom_embedding,
    chgnet.bond_embedding,
    chgnet.angle_embedding,
    chgnet.bond_basis_expansion,
    chgnet.angle_basis_expansion,
    chgnet.atom_conv_layers[:-1],
    chgnet.bond_conv_layers,
    chgnet.angle_layers,
]:
    for param in layer.parameters():
        param.requires_grad = False

# Define Trainer
trainer = Trainer(
    model=chgnet,
    targets="efsm",
    optimizer="Adam",
    scheduler="CosLR",
    criterion="MSE",
    epochs=500,
    learning_rate=1e-4,
    use_device="cuda",
    print_freq=6,
)

After tunning, the accuracy seems great.

Epoch: [499][1/9] | Time (0.334)(0.000) | Loss 0.0533(0.0533) | MAE e 0.006(0.006)  f 0.147(0.147)  s 0.214(0.214)  m 0.118(0.118)  
Epoch: [499][6/9] | Time (0.332)(0.000) | Loss 0.0504(0.0490) | MAE e 0.006(0.006)  f 0.147(0.142)  s 0.156(0.173)  m 0.118(0.119)  
*   e_MAE (0.009) 	f_MAE (0.161) 	s_MAE (0.313) 	m_MAE (0.122) 	
---------Evaluate Model on Test Set---------------
**  e_MAE (0.004) 	f_MAE (0.161) 	s_MAE (0.219) 	m_MAE (0.123)  

However, when I try to load the fine-tuned model for some validation, something went wrong. Firstly, I load the original model and fine-tuned model by:

chgnet = CHGNet.load()
chgnet_ft = CHGNet.from_file("01-30-2025/bestE_epoch499_e9_f161_s313_m122.pth.tar")

MLFF = CHGNetCalculator(
        potential=chgnet,
        properties = ['energy', 'forces'],
        use_device='cuda')

MLFF_FT = CHGNetCalculator(
        potential=chgnet_ft,
        properties = ['energy', 'forces'],
        use_device='cuda')

Then when I applied these two model as ASE calculator to calculate the force in the snapshots strcuture, I found that the output of both model is actually the same.

Does this means that the fine-tuned model is not loaded correctly, or the size of my dataset is to small to have a positive effect in fine-tuning?

[M9JS]

Problem solved now, I found I make a mistake in using CHGNetCalculator.

chgnet = CHGNet.load()
chgnet_ft = CHGNet.from_file("01-30-2025/bestE_epoch499_e9_f161_s313_m122.pth.tar")

MLFF = CHGNetCalculator(
        potential=chgnet,
        properties = ['energy', 'forces'],
        use_device='cuda')

MLFF_FT = CHGNetCalculator(
        potential=chgnet_ft,
        properties = ['energy', 'forces'],
        use_device='cuda')

It should be
MLFF = CHGNetCalculator(
model=chgnet_ft,
properties = ['energy', 'forces'],
use_device='cuda')