Add TS test for comples Spectrogram

pytorch · Apr 1, 2021 · 7ff9741 · 7ff9741
1 parent b395211
commit 7ff9741
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/test/torchaudio_unittest/transforms/torchscript_consistency_impl.py b/test/torchaudio_unittest/transforms/torchscript_consistency_impl.py
@@ -21,6 +21,10 @@ def test_Spectrogram(self):
         tensor = torch.rand((1, 1000))
         self._assert_consistency(T.Spectrogram(), tensor)
 
+    def test_Spectrogram_return_complex(self):
+        tensor = torch.rand((1, 1000))
+        self._assert_consistency(T.Spectrogram(power=None, return_complex=True), tensor)
+
     def test_GriffinLim(self):
         tensor = torch.rand((1, 201, 6))
         self._assert_consistency(T.GriffinLim(length=1000, rand_init=False), tensor)