From cb0040b508f29ccc55f68ff936dbe1cf863a3cba Mon Sep 17 00:00:00 2001
From: shahules786 <shahules786@gmail.com>
Date: Wed, 12 Oct 2022 18:44:22 +0530
Subject: [PATCH] rename loss

---
 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 3ad5fa7..4f055b4 100644
--- a/enhancer/models/model.py
+++ b/enhancer/models/model.py
@@ -135,7 +135,7 @@ class Model(pl.LightningModule):
         loss = self.loss(prediction, target)
 
         self.log(
-            f"train_{self.loss.name}",
+            "train_loss",
             loss.item(),
             on_epoch=True,
             on_step=True,
@@ -152,6 +152,7 @@ class Model(pl.LightningModule):
         target = batch["clean"]
         prediction = self(mixed_waveform)
 
+        metric_dict["valid_loss"] = self.loss(target, prediction).item()
         for metric in self.metric:
             value = metric(target, prediction)
             metric_dict[f"valid_{metric.name}"] = value.item()