diff --git a/enhancer/loss.py b/enhancer/loss.py index 693e610..d673ec3 100644 --- a/enhancer/loss.py +++ b/enhancer/loss.py @@ -1,3 +1,5 @@ +from modulefinder import Module +from turtle import forward import torch import torch.nn as nn @@ -32,6 +34,20 @@ class mean_absolute_error(nn.Module): return self.loss_fun(prediction, target) +class Si_SDR(nn.Module): + + def __init__( + self + ): + pass + + 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""") + prediction,target = prediction.unsqueeze(1),target.unsqueeze(1) + class Avergeloss(nn.Module): def __init__(self,losses):