Merge 72ad287 into ae4dca9

Author: awaelchli

pytorch_lightning/accelerators/gpu.py

@@ -27,6 +27,7 @@ def on_train_start(self):
 
     def on_train_end(self):
         # clean up memory
+        self.model.cpu()
         with torch.cuda.device(self.root_device):
             torch.cuda.empty_cache()
 

pytorch_lightning/trainer/training_loop.py

@@ -148,13 +148,7 @@ def on_train_end(self):
             self.trainer.profiler.describe()
 
         # give accelerators a chance to finish
-        self.trainer.accelerator_backend.on_train_end()
-
-        # clear mem
-        if self.trainer._device_type == DeviceType.GPU:
-            model = self.trainer.get_model()
-            model.cpu()
-            torch.cuda.empty_cache()
+        self.trainer.accelerator.on_train_end()
 
     def check_checkpoint_callback(self, should_update, is_last=False):
         # TODO bake this logic into the ModelCheckpoint callback