cheat

Author: awaelchli

pytorch_lightning/trainer/trainer.py

@@ -992,6 +992,7 @@ def _run_evaluation(self) -> _EVALUATE_OUTPUT:
                 del batch
                 import gc
                 gc.collect()
+                torch.cuda.empty_cache()
 
                 # log batch metrics
                 self.logger_connector.update_eval_step_metrics()
@@ -1102,6 +1103,7 @@ def _run_predict(self) -> Optional[_PREDICT_OUTPUT]:
                 del batch
                 import gc
                 gc.collect()
+                torch.cuda.empty_cache()
 
         # call hook
         results = self.predict_loop.on_predict_epoch_end()

pytorch_lightning/trainer/training_loop.py

@@ -485,6 +485,7 @@ def run_training_epoch(self):
             del batch
             import gc
             gc.collect()
+            torch.cuda.empty_cache()
 
             # -----------------------------------------
             # SAVE METRICS TO LOGGERS AND PROGRESS_BAR