hack

Author: awaelchli

pytorch_lightning/trainer/trainer.py

@@ -990,6 +990,8 @@ def _run_evaluation(self) -> _EVALUATE_OUTPUT:
 
                 # release memory before running any other hooks
                 del batch
+                import gc
+                gc.collect()
 
                 # log batch metrics
                 self.logger_connector.log_evaluation_step_metrics()
@@ -1101,6 +1103,8 @@ def _run_predict(self) -> Optional[_PREDICT_OUTPUT]:
 
                 # release memory before running any other hooks
                 del batch
+                import gc
+                gc.collect()
 
         # call hook
         results = self.predict_loop.on_predict_epoch_end()

pytorch_lightning/trainer/training_loop.py

@@ -510,6 +510,8 @@ def run_training_epoch(self):
 
             # release memory before running any other hooks
             del batch
+            import gc
+            gc.collect()
 
             # -----------------------------------------
             # SAVE METRICS TO LOGGERS