diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 818c1947..1e47e3e7 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -144,7 +144,7 @@ class TrainerIO(object): filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number) # give model a chance to do something on hpc_save - self.on_hpc_save() + self.model.on_hpc_save() # request what to save from the model checkpoint_dict = self.dump_checkpoint() @@ -168,7 +168,7 @@ class TrainerIO(object): model.load_model_specific(checkpoint) # call model hook - self.on_hpc_load() + self.model.on_hpc_load() def max_ckpt_in_folder(self, path): files = os.listdir(path)