update

toolbox/torchaudio/losses/spectral.py

@@ -158,9 +158,11 @@ class SpectralConvergenceLoss(torch.nn.Module):
 
     def __init__(self,
                  reduction: str = "mean",
+                 eps: float = 1e-8,
                  ):
         super(SpectralConvergenceLoss, self).__init__()
         self.reduction = reduction
+        self.eps = eps
 
         if reduction not in ("sum", "mean"):
             raise AssertionError(f"param reduction must be sum or mean.")
@@ -175,19 +177,9 @@ class SpectralConvergenceLoss(torch.nn.Module):
         :return:
         """
         error_norm = torch.norm(denoise_magnitude - clean_magnitude, p="fro", dim=(-1, -2))
-        if torch.any(torch.isnan(error_norm)) or torch.any(torch.isinf(error_norm)):
-            raise AssertionError("SpectralConvergenceLoss, nan or inf in error_norm")
         truth_norm = torch.norm(clean_magnitude, p="fro", dim=(-1, -2))
-        if torch.any(torch.isnan(truth_norm)):
-            raise AssertionError("SpectralConvergenceLoss, nan in truth_norm")
-        if torch.any(torch.isinf(truth_norm)):
-            raise AssertionError("SpectralConvergenceLoss, inf in truth_norm")
 
-        batch_loss = error_norm / truth_norm
-        if torch.any(torch.isnan(batch_loss)):
-            raise AssertionError("SpectralConvergenceLoss, nan in batch_loss")
-        if torch.any(torch.isinf(batch_loss)):
-            raise AssertionError("SpectralConvergenceLoss, inf in batch_loss")
+        batch_loss = error_norm / (truth_norm + self.eps)
 
         if self.reduction == "mean":
             loss = torch.mean(batch_loss)
@@ -196,9 +188,6 @@ class SpectralConvergenceLoss(torch.nn.Module):
         else:
             raise AssertionError
 
-        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
-            raise AssertionError("SpectralConvergenceLoss, nan or inf in loss")
-
         return loss