update tests and TimeStretch

pytorch · Jul 1, 2020 · 6795dc4 · 6795dc4
1 parent 2f6494b
commit 6795dc4
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/test/test_batch_consistency.py b/test/test_batch_consistency.py
@@ -246,6 +246,7 @@ def test_batch_TimeStretch(self):
         rate = 2
 
         complex_specgrams = torch.stft(waveform, **kwargs)
+        complex_specgrams = torch.view_as_complex(complex_specgrams)
 
         # Single then transform then batch
         expected = torchaudio.transforms.TimeStretch(

diff --git a/test/torchscript_consistency_impl.py b/test/torchscript_consistency_impl.py
@@ -586,7 +586,7 @@ def test_TimeStretch(self):
         n_freq = 400
         hop_length = 512
         fixed_rate = 1.3
-        tensor = torch.rand((10, 2, n_freq, 10, 2))
+        tensor = torch.rand((10, 2, n_freq, 10), dtype=torch.cdouble)
         self._assert_consistency(
             T.TimeStretch(n_freq=n_freq, hop_length=hop_length, fixed_rate=fixed_rate),
             tensor,

diff --git a/torchaudio/transforms.py b/torchaudio/transforms.py
@@ -686,7 +686,7 @@ def forward(self, complex_specgrams: Tensor, overriding_rate: Optional[float] =
         Returns:
             Tensor: Stretched complex spectrogram of dimension (..., freq, ceil(time/rate), complex=2).
         """
-        assert complex_specgrams.size(-1) == 2, "complex_specgrams should be a complex tensor, shape (..., complex=2)"
+        # assert complex_specgrams.size(-1) == 2, "complex_specgrams should be a complex tensor, shape (..., complex=2)"
 
         if overriding_rate is None:
             rate = self.fixed_rate