autocast

Lightning-AI · Apr 22, 2020 · 2763a48 · 2763a48
1 parent 801d27d
commit 2763a48
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py
@@ -530,6 +530,10 @@ def dp_train(self, model):
 
         model.cuda(self.root_gpu)
 
+        # hack forward to do autocast for the user
+        if self.use_amp and self.use_native_amp:
+            model.forward = torch.cuda.amp.autocast()(model.forward)
+
         # TODO: remove in v0.8.0
         # check for this bug (amp + dp + !01 doesn't work)
         # https://github.com/NVIDIA/apex/issues/227