From c5fdfbe188ad42117af4b9567f4f1dd1a6faed40 Mon Sep 17 00:00:00 2001
From: shahules786 <shahules786@gmail.com>
Date: Wed, 14 Sep 2022 11:50:30 +0530
Subject: [PATCH] shape/value checks

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

diff --git a/enhancer/models/demucs.py b/enhancer/models/demucs.py
index 2ac067b..d4e35a9 100644
--- a/enhancer/models/demucs.py
+++ b/enhancer/models/demucs.py
@@ -100,6 +100,9 @@ class Demucs(Model):
         if mixed_signal.dim() == 2:
             mixed_signal = mixed_signal.unsqueeze(1)
 
+        if mixed_signal.size(1)!=1:
+            raise TypeError(f"Demucs can only process mono channel audio, input has {mixed_signal.size(1)} channels")
+
         length = mixed_signal.shape[-1]
         x = F.pad(mixed_signal, (0,self.get_padding_length(length) - length)) 
         if self.hparams.resample>1:
diff --git a/enhancer/models/model.py b/enhancer/models/model.py
index 2adfef8..275027b 100644
--- a/enhancer/models/model.py
+++ b/enhancer/models/model.py
@@ -3,7 +3,7 @@ from torch.optim import Adam
 import pytorch_lightning as pl
 
 from enhancer.data.dataset import Dataset
-from enhancer.utils.loss import Avergeloss
+from enhancer.utils.loss import LOSS_MAP, Avergeloss
 
 
 class Model(pl.LightningModule):