Seeming disparity in calling `on_load_checkpoint` in hpc load/save

[ivannz]

I noticed that before calling model.on_hpc_save(checkpoint) on L464, .hpc_save() gives the model a chance to add a few things by calling .dump_checkpoint() which later calls model.on_save_checkpoint(checkpoint) on L369.

However .hpc_load() calls only model.on_hpc_load(checkpoint) at training_io.py#L502, and does not seem to call model.on_load_checkpoint neither via .restore_training_state(), nor directly. This seems to be in contrast to what .restore() does on L305.

Is the intended behaviour?

[github-actions]

Hi! thanks for your contribution!, great first issue!

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.