diff --git a/enhancer/models/model.py b/enhancer/models/model.py
index 4f1d75b..7dbc065 100644
--- a/enhancer/models/model.py
+++ b/enhancer/models/model.py
@@ -110,6 +110,7 @@ class Model(pl.LightningModule):
 
     def setup(self, stage: Optional[str] = None):
         if stage == "fit":
+            torch.cuda.empty_cache()
             self.dataset.setup(stage)
             self.dataset.model = self