How to pass the composition model to CHGNet while training?

[naveenmohandas]

Hi,
I was using a dataset not developed based on the MPRelaxSet parameters. It is mentioned that the composition model should be used to overcome the energy shifts. However, I couldn't find how to use the composition model in the example.
Is this the right way to pass the composition model while training?

On a side note, I also couldn't understand how the model fits to the energies from the dataset. Could you briefly explain how it works?

new_atom_ref = AtomRef(is_intensive=True)
new_atom_ref.initialize_from_MPtrj()
new_atom_ref.fit(structures, energies_per_atom)

model = CHGNet.load()
model.composition_model = new_atom_ref