Fix amp restore place

Author: tarepan

pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -127,12 +127,6 @@ def restore_model_state(self, model: LightningModule, checkpoint) -> None:
         # restore the state_dict on the model
         model.load_state_dict(checkpoint['state_dict'])
 
-        # restore amp scaling
-        if self.trainer.amp_backend == AMPType.NATIVE and 'native_amp_scaling_state' in checkpoint:
-            self.trainer.scaler.load_state_dict(checkpoint['native_amp_scaling_state'])
-        elif self.trainer.amp_backend == AMPType.APEX and 'amp_scaling_state' in checkpoint:
-            amp.load_state_dict(checkpoint['amp_scaling_state'])
-
     def restore_training_state(self, checkpoint):
         """
         Restore trainer state.
@@ -155,6 +149,12 @@ def restore_training_state(self, checkpoint):
                 " where `model.ckpt` is your checkpoint file."
             )
 
+        # restore amp scaling
+        if self.trainer.amp_backend == AMPType.NATIVE and 'native_amp_scaling_state' in checkpoint:
+            self.trainer.scaler.load_state_dict(checkpoint['native_amp_scaling_state'])
+        elif self.trainer.amp_backend == AMPType.APEX and 'amp_scaling_state' in checkpoint:
+            amp.load_state_dict(checkpoint['amp_scaling_state'])
+
         # restore callback states
         self.trainer.on_load_checkpoint(checkpoint)