update

toolbox/torchaudio/models/clean_unet/loss.py

@@ -4,12 +4,6 @@ import torch
 import torch
 import torch.nn.functional as F
 
-# from distutils.version import LooseVersion
-
-
-# is_pytorch_17plus = LooseVersion(torch.__version__) >= LooseVersion("1.7")
-is_pytorch_17plus = True
-
 
 def stft(x, fft_size, hop_size, win_length, window):
     """
@@ -22,18 +16,9 @@ def stft(x, fft_size, hop_size, win_length, window):
     :return: Magnitude spectrogram (B, #frames, fft_size // 2 + 1).
     """
 
-    if is_pytorch_17plus:
-        x_stft = torch.stft(
-            x, fft_size, hop_size, win_length, window, return_complex=False
-        )
-    else:
-        x_stft = torch.stft(x, fft_size, hop_size, win_length, window)
-    real = x_stft[..., 0]
-    imag = x_stft[..., 1]
-
-    # NOTE(kan-bayashi): clamp is needed to avoid nan or inf
-    return torch.sqrt(torch.clamp(real**2 + imag**2, min=1e-7)).transpose(2, 1)
+    x_stft = torch.stft(x, fft_size, hop_size, win_length, window)
 
+    return x_stft.abs()
 
 class SpectralConvergenceLoss(torch.nn.Module):
     """Spectral convergence loss module."""