Image training example save_model without loss_scaler uses save_checkpoint

Hello,
In the load_and_save.py file of the image example, the save_model function uses model.save_checkpoint but I don't think any of the model in the image example has a save_checkpoint method.
There is a save_checkpoint method in the state.py file under examples/text/logic/ but it is for the TrainState class which isn't imported here.

Maybe the save_on_master function can also be used in cases where there is no loss_scaler?

Maybe I missed something, in which case I apologize.