adopt changes from pytorch#1009

mthrok · Nov 6, 2020 · 55e0c77 · 55e0c77
1 parent 0b89f9a
commit 55e0c77
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/torchaudio/functional.py b/torchaudio/functional.py
@@ -104,6 +104,12 @@ def spectrogram(
     if normalized:
         spec_f /= window.pow(2.).sum().sqrt()
     if power is not None:
+        if power == 1.0:
+            return spec_f.abs()
+        if power == 2.0:
+            if spec_f.is_cuda:
+                return torch.view_as_real(spec_f).pow(power).sum(-1)
+            return spec_f.pow(power)
         return spec_f.abs().pow(power)
     return torch.view_as_real(spec_f)