Please add enhancement feature so that tool supports with model.fit()

Hello,

Please add enhancement feature such tool supports with model.fit() for any number of epochs for tracking the energy and prediction of energy.

Thanks and Regards,
Prince Chaturvedi