From c193a48e8e27a2779273bb28d7600ef133fea858 Mon Sep 17 00:00:00 2001
From: shahules786 <shahules786@gmail.com>
Date: Mon, 10 Oct 2022 21:03:39 +0530
Subject: [PATCH] empty cache

---
 enhancer/models/model.py | 1 +
 1 file changed, 1 insertion(+)

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