Move result teardown to loops

pytorch_lightning/loops/dataloader/evaluation_loop.py

@@ -266,3 +266,7 @@ def on_evaluation_epoch_end(self) -> None:
         self.trainer.call_hook(hook_name)
         self.trainer.call_hook("on_epoch_end")
         self.trainer.logger_connector.on_epoch_end()
+
+    def teardown(self) -> None:
+        self._results.cpu()
+        self.epoch_loop.teardown()

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -222,7 +222,10 @@ def on_run_end(self) -> List[List[STEP_OUTPUT]]:
 
     def teardown(self) -> None:
         """Frees memory of tracked epoch outputs."""
-        self.epoch_output = None
+        self._epoch_output = None
+        self._results.cpu()
+        self.batch_loop.teardown()
+        self.val_loop.teardown()
 
     def _run_validation(self):
         # reload dataloaders

pytorch_lightning/loops/fit_loop.py

@@ -286,3 +286,6 @@ def state_dict(self) -> Dict:
 
     def load_state_dict(self, state_dict: Dict) -> None:
         self.epoch_loop.load_state_dict(state_dict["epoch_loop"])
+
+    def teardown(self) -> None:
+        self.epoch_loop.teardown()

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -312,9 +312,3 @@ def progress_bar_metrics(self) -> Dict[str, float]:
             metrics = self.metrics[MetricSource.PBAR]
             self._progress_bar_metrics.update(metrics)
         return self._progress_bar_metrics
-
-    def teardown(self):
-        self.trainer.fit_loop.epoch_loop._results.cpu()
-        self.trainer.fit_loop.epoch_loop.val_loop._results.cpu()
-        self.trainer.validate_loop._results.cpu()
-        self.trainer.test_loop._results.cpu()

pytorch_lightning/trainer/trainer.py

@@ -901,7 +901,6 @@ def _pre_dispatch(self):
     def _post_dispatch(self):
         self.accelerator.post_dispatch(self)
         self.accelerator.teardown()
-        self.logger_connector.teardown()
 
     def _dispatch(self):
         if self.evaluating: