Call on_load_checkpoint before loading state_dict (#4057)

Lightning-AI · Oct 14, 2020 · dec31b3 · dec31b3
1 parent f967fbb
commit dec31b3
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/pytorch_lightning/trainer/connectors/checkpoint_connector.py b/pytorch_lightning/trainer/connectors/checkpoint_connector.py
@@ -101,15 +101,16 @@ def restore(self, checkpoint_path: str, on_gpu: bool):
         # load model state
         model = self.trainer.get_model()
 
-        # load the state_dict on the model automatically
-        model.load_state_dict(checkpoint['state_dict'])
-
         # give the datamodule a chance to load something
         if self.trainer.datamodule is not None:
             self.trainer.datamodule.on_load_checkpoint(checkpoint)
+
         # give model a chance to load something
         model.on_load_checkpoint(checkpoint)
 
+        # load the state_dict on the model automatically
+        model.load_state_dict(checkpoint['state_dict'])
+
         if on_gpu:
             model.cuda(self.trainer.root_gpu)