Make power and normalized configurable for torchaudio.transforms.MelS…

…pectrogram (#633)
pytorch · May 13, 2020 · 36a1e90 · 36a1e90
1 parent b7f192f
commit 36a1e90
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/torchaudio/transforms.py b/torchaudio/transforms.py
@@ -395,20 +395,24 @@ def __init__(self,
                  pad: int = 0,
                  n_mels: int = 128,
                  window_fn: Callable[..., Tensor] = torch.hann_window,
+                 power: Optional[float] = 2.,
+                 normalized: bool = False,
                  wkwargs: Optional[dict] = None) -> None:
         super(MelSpectrogram, self).__init__()
         self.sample_rate = sample_rate
         self.n_fft = n_fft
         self.win_length = win_length if win_length is not None else n_fft
         self.hop_length = hop_length if hop_length is not None else self.win_length // 2
         self.pad = pad
+        self.power = power
+        self.normalized = normalized
         self.n_mels = n_mels  # number of mel frequency bins
         self.f_max = f_max
         self.f_min = f_min
         self.spectrogram = Spectrogram(n_fft=self.n_fft, win_length=self.win_length,
                                        hop_length=self.hop_length,
-                                       pad=self.pad, window_fn=window_fn, power=2.,
-                                       normalized=False, wkwargs=wkwargs)
+                                       pad=self.pad, window_fn=window_fn, power=self.power,
+                                       normalized=self.normalized, wkwargs=wkwargs)
         self.mel_scale = MelScale(self.n_mels, self.sample_rate, self.f_min, self.f_max, self.n_fft // 2 + 1)
 
     def forward(self, waveform: Tensor) -> Tensor: