From d99fd0eb61ff4d969f15ac4a05db0c209c209dc9 Mon Sep 17 00:00:00 2001
From: shahules786 <shahules786@gmail.com>
Date: Sun, 16 Oct 2022 11:13:44 +0530
Subject: [PATCH] fix duration estimation

---
 enhancer/models/model.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/enhancer/models/model.py b/enhancer/models/model.py
index aa0199b..dc5219d 100644
--- a/enhancer/models/model.py
+++ b/enhancer/models/model.py
@@ -115,17 +115,23 @@ class Model(pl.LightningModule):
             self.dataset.setup(stage)
             print(
                 "Total train duration",
-                self.dataset.train_dataloader().dataset.__len__() / 60,
+                self.dataset.train_dataloader().dataset.__len__()
+                * self.dataset.duration
+                / 60,
                 "minutes",
             )
             print(
                 "Total validation duration",
-                self.dataset.val_dataloader().dataset.__len__() / 60,
+                self.dataset.val_dataloader().dataset.__len__()
+                * self.dataset.duration
+                / 60,
                 "minutes",
             )
             print(
                 "Total test duration",
-                self.dataset.test_dataloader().dataset.__len__() / 60,
+                self.dataset.test_dataloader().dataset.__len__()
+                * self.dataset.duration
+                / 60,
                 "minutes",
             )
             self.dataset.model = self