Add garbage collection to predict_step

project-lighter · Feb 2, 2024 · 23fa0e9 · 23fa0e9
1 parent b95775c
commit 23fa0e9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/lighter/system.py b/lighter/system.py
@@ -2,6 +2,7 @@
 
 from functools import partial
 
+import gc
 import pytorch_lightning as pl
 import torch
 from loguru import logger
@@ -183,6 +184,7 @@ def _base_step(self, batch: Union[List, Tuple], batch_idx: int, mode: str) -> Un
         if mode == "predict":
             # Postprocessing for logging/writing.
             pred = apply_fns(pred, self.postprocessing["logging"]["pred"])
+            gc.collect()
             return {"pred": pred, "id": id}
 
         # Calculate the loss.