diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py index 03cb5864a8..142d2b337b 100644 --- a/pytorch_lightning/root_module/model_saving.py +++ b/pytorch_lightning/root_module/model_saving.py @@ -79,7 +79,7 @@ class TrainerIO(object): # add the state_dict from the model model = self.__get_model() - checkpoint['state_dict'] = model.get_state_dict + checkpoint['state_dict'] = model.state_dict() # give the model a chance to add a few things model.on_save_checkpoint(checkpoint)