autocast

Lightning-AI · Apr 22, 2020 · 0a20aee · 0a20aee
1 parent f0d2c78
commit 0a20aee
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/pytorch_lightning/trainer/distrib_parts.py b/pytorch_lightning/trainer/distrib_parts.py
@@ -520,7 +520,9 @@ def dp_train(self, model):
         model.cuda(self.root_gpu)
 
         # hack forward to do autocast for the user
+        model_autocast_original_forward = model.forward
         if self.use_amp and self.use_native_amp:
+            # wrap the user's forward in autocast and give it back at the end
             model.forward = torch.cuda.amp.autocast()(model.forward)
 
         # TODO: remove in v0.8.0
@@ -547,6 +549,9 @@ def dp_train(self, model):
 
         self.run_pretrain_routine(model)
 
+        # when training completes give back the forward
+        model.forward = model_autocast_original_forward
+
 
 def normalize_parse_gpu_string_input(s):
     if isinstance(s, str):

diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
@@ -489,6 +489,8 @@ def __init__(
 
         # AMP init
         # These are the only lines needed after v0.8.0
+        # we wrap the user's forward with autocast and give it back at the end of fit
+        self.autocast_original_forward = None
         self.use_native_amp = hasattr(torch.cuda, "amp") and hasattr(torch.cuda.amp, "autocast")
         if self.use_native_amp and self.precision == 16:
             self.scaler = torch.cuda.amp.GradScaler()