Seeming disparity in calling `on_load_checkpoint` in hpc load/save
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?