fixed native amp + ddp (#1788)

* fixed native amp + ddp * fixed native amp + ddp
Lightning-AI · May 12, 2020 · 7b60d49 · 7b60d49
1 parent 1df0d2d
commit 7b60d49
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
@@ -510,8 +510,7 @@ def __init__(
         self.autocast_original_forward = None
         self.use_native_amp = hasattr(torch.cuda, "amp") and hasattr(torch.cuda.amp, "autocast")
         self.precision = precision
-        if self.use_native_amp and self.precision == 16:
-            self.scaler = torch.cuda.amp.GradScaler()
+        self.scaler = None
 
         # TODO: remove for v0.8.0
         self.amp_level = amp_level
@@ -858,6 +857,10 @@ def run_pretrain_routine(self, model: LightningModule):
         # set local properties on the model
         self.copy_trainer_model_properties(ref_model)
 
+        # init amp. Must be done here instead of __init__ to allow ddp to work
+        if self.use_native_amp and self.precision == 16:
+            self.scaler = torch.cuda.amp.GradScaler()
+
         # log hyper-parameters
         if self.logger is not None:
             # save exp to get started