From e4cd6ef10e13080f71caf53bf4dda1994982f8d1 Mon Sep 17 00:00:00 2001 From: shahules786 Date: Tue, 13 Sep 2022 22:02:09 +0530 Subject: [PATCH] add typeerror in shape mismatch --- enhancer/utils/loss.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/enhancer/utils/loss.py b/enhancer/utils/loss.py index a51669c..7e555b5 100644 --- a/enhancer/utils/loss.py +++ b/enhancer/utils/loss.py @@ -11,6 +11,10 @@ class mean_squared_error(nn.Module): self.loss_fun = nn.MSELoss(reduction=reduction) def forward(self,prediction:torch.Tensor, target: torch.Tensor): + + if prediction.size() != target.size() or target.ndim < 3: + raise TypeError(f"""Inputs must be of the same shape (batch_size,channels,samples) + got {prediction.size()} and {target.size()} instead""") return self.loss_fun(prediction, target) @@ -23,6 +27,10 @@ class mean_absolute_error(nn.Module): def forward(self, prediction:torch.Tensor, target: torch.Tensor): + if prediction.size() != target.size() or target.ndim < 3: + raise TypeError(f"""Inputs must be of the same shape (batch_size,channels,samples) + got {prediction.size()} and {target.size()} instead""") + return self.loss_fun(prediction, target) class Avergeloss(nn.Module):