move device-specific teardown logic from training loop to accelerator (…

…#5973) * on train end * switch order
Lightning-AI · Feb 15, 2021 · aa60c08 · aa60c08
1 parent ae4dca9
commit aa60c08
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 7 deletions.
diff --git a/pytorch_lightning/accelerators/gpu.py b/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()
 

diff --git a/pytorch_lightning/trainer/training_loop.py b/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