update

toolbox/torchaudio/models/nx_clean_unet/discriminator.py

@@ -22,15 +22,16 @@ class MetricDiscriminator(nn.Module):
         self.win_length = config.win_length
         self.hop_length = config.hop_length
 
+        self.transform = torchaudio.transforms.Spectrogram(
+            n_fft=self.n_fft,
+            win_length=self.win_length,
+            hop_length=self.hop_length,
+            power=1.0,
+            window_fn=torch.hann_window,
+            # window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
+        )
+
         self.layers = nn.Sequential(
-            torchaudio.transforms.Spectrogram(
-                n_fft=self.n_fft,
-                win_length=self.win_length,
-                hop_length=self.hop_length,
-                power=1.0,
-                window_fn=torch.hamming_window,
-                # window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
-            ),
             nn.utils.spectral_norm(nn.Conv2d(self.in_channel, dim, (4,4), (2,2), (1,1), bias=False)),
             nn.InstanceNorm2d(dim, affine=True),
             nn.PReLU(dim),
@@ -53,6 +54,9 @@ class MetricDiscriminator(nn.Module):
         )
 
     def forward(self, x, y):
+        x = self.transform.forward(x)
+        y = self.transform.forward(y)
+
         xy = torch.stack((x, y), dim=1)
         return self.layers(xy)