si-sdr

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):