Add autograd test for T.GriffinLim (#1421)

pytorch · Apr 6, 2021 · 6929331 · 6929331
1 parent b388d48
commit 6929331
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/test/torchaudio_unittest/transforms/autograd_test_impl.py b/test/torchaudio_unittest/transforms/autograd_test_impl.py
@@ -65,6 +65,15 @@ def test_melspectrogram(self):
         waveform = get_whitenoise(sample_rate=sample_rate, duration=0.05, n_channels=2)
         self.assert_grad(transform, [waveform], nondet_tol=1e-10)
 
+    @parameterized.expand([(0, ), (0.99, )])
+    def test_griffinlim(self, momentum):
+        n_fft = 400
+        n_frames = 5
+        n_iter = 3
+        spec = torch.rand(n_fft // 2 + 1, n_frames) * n_fft
+        transform = T.GriffinLim(n_fft=n_fft, n_iter=n_iter, momentum=momentum, rand_init=False)
+        self.assert_grad(transform, [spec], nondet_tol=1e-10)
+
     @parameterized.expand([(False, ), (True, )])
     def test_mfcc(self, log_mels):
         sample_rate = 8000

diff --git a/torchaudio/functional/functional.py b/torchaudio/functional/functional.py
@@ -203,7 +203,7 @@ def griffinlim(
                               hop_length=hop_length,
                               win_length=win_length,
                               window=window,
-                              length=length).float()
+                              length=length)
 
         # Rebuild the spectrogram
         rebuilt = torch.view_as_real(

diff --git a/torchaudio/transforms.py b/torchaudio/transforms.py
@@ -172,7 +172,7 @@ def __init__(self,
         super(GriffinLim, self).__init__()
 
         assert momentum < 1, 'momentum={} > 1 can be unstable'.format(momentum)
-        assert momentum > 0, 'momentum={} < 0'.format(momentum)
+        assert momentum >= 0, 'momentum={} < 0'.format(momentum)
 
         self.n_fft = n_fft
         self.n_iter = n_iter