From ea5c78798add582922f36fd6bbc38194776b13f8 Mon Sep 17 00:00:00 2001
From: shahules786 <shahules786@gmail.com>
Date: Sun, 23 Oct 2022 12:33:38 +0530
Subject: [PATCH] model assigment'

---
 enhancer/models/model.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/enhancer/models/model.py b/enhancer/models/model.py
index dc5219d..3b60b85 100644
--- a/enhancer/models/model.py
+++ b/enhancer/models/model.py
@@ -113,6 +113,8 @@ class Model(pl.LightningModule):
         if stage == "fit":
             torch.cuda.empty_cache()
             self.dataset.setup(stage)
+            self.dataset.model = self
+
             print(
                 "Total train duration",
                 self.dataset.train_dataloader().dataset.__len__()
@@ -134,7 +136,6 @@ class Model(pl.LightningModule):
                 / 60,
                 "minutes",
             )
-            self.dataset.model = self
 
     def train_dataloader(self):
         return self.dataset.train_dataloader()