jit for fbank & specify spectral_magnitude args

anautsch · anautsch · commit 5fe6ff573f8a · 2023-03-21T14:35:19.000+01:00
diff --git a/speechbrain/lobes/features.py b/speechbrain/lobes/features.py
@@ -77,6 +77,8 @@ class Fbank(torch.nn.Module):
     >>> feats = feature_maker(inputs)
     >>> feats.shape
     torch.Size([10, 101, 40])
+    >>> fbank = Fbank(sample_rate=8000, n_fft=512, n_mels=64, left_frames=0, right_frames=0, deltas=False)
+    >>> _ = torch.jit.script(fbank)
     """
 
     def __init__(
diff --git a/speechbrain/processing/features.py b/speechbrain/processing/features.py
@@ -310,7 +310,9 @@ def forward(self, x, sig_length=None):
         return istft
 
 
-def spectral_magnitude(stft, power=1, log=False, eps=1e-14):
+def spectral_magnitude(
+    stft, power: int = 1, log: bool = False, eps: float = 1e-14
+):
     """Returns the magnitude of a complex spectrogram.
 
     Arguments
