Add torchscript test to oscillator_bank (#2864)

Summary:
Missing from https://github.com/pytorch/audio/issues/2848

Pull Request resolved: https://github.com/pytorch/audio/pull/2864

Reviewed By: carolineechen

Differential Revision: D41413381

Pulled By: mthrok

fbshipit-source-id: 4377ed4a59504c6ade9ee6f42938a2bc3f04fb73

test/torchaudio_unittest/prototype/functional/torchscript_consistency_test_impl.py

@@ -58,3 +58,10 @@ class TorchScriptConsistencyTestImpl(TestBaseMixin):
         n_barks = 10
         sample_rate = 16000
         self._assert_consistency(F.barkscale_fbanks, (n_stft, f_min, f_max, n_barks, sample_rate, "traunmuller"))
+
+    def test_oscillator_bank(self):
+        num_frames, num_pitches, sample_rate = 8000, 8, 8000
+        freq = torch.rand((num_frames, num_pitches), dtype=self.dtype, device=self.device)
+        amps = torch.ones_like(freq)
+
+        self._assert_consistency(F.oscillator_bank, (freq, amps, sample_rate, "sum"))

torchaudio/prototype/functional/_dsp.py

@@ -71,7 +71,7 @@ def oscillator_bank(
     # We might add angular_cumsum if it turned out to be undesirable.
     pi2 = 2.0 * torch.pi
     freqs = frequencies * pi2 / sample_rate % pi2
-    phases = torch.cumsum(freqs, axis=-2)
+    phases = torch.cumsum(freqs, dim=-2)
 
     waveform = amplitudes * torch.sin(phases)
     if reduction == "sum":