From 042ce2b525db8095a5595d2c7605228c20fdfa2d Mon Sep 17 00:00:00 2001 From: shahules786 Date: Wed, 14 Sep 2022 11:51:03 +0530 Subject: [PATCH] add valuerror --- enhancer/utils/loss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/enhancer/utils/loss.py b/enhancer/utils/loss.py index 7e555b5..3c1cc06 100644 --- a/enhancer/utils/loss.py +++ b/enhancer/utils/loss.py @@ -11,7 +11,7 @@ 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""") @@ -46,7 +46,7 @@ class Avergeloss(nn.Module): def validate_loss(self,loss:str): if loss not in LOSS_MAP.keys(): - raise ValueError() + raise ValueError(f"Invalid loss function {loss}, available loss functions are {LOSS_MAP.keys()}") else: return LOSS_MAP[loss]