Skip to content

Commit 89ee52e

Browse files
committed
Add dispatcher for mac core loop
1 parent 69c7ecc commit 89ee52e

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

src/torchaudio/functional/filtering.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -931,9 +931,14 @@ def _lfilter_core_generic_loop(input_signal_windows: Tensor, a_coeffs_flipped: T
931931
o0 -= (windowed_output_signal.transpose(0, 1) @ a_coeffs_flipped)[..., 0].t()
932932
padded_output_waveform[:, :, i_sample + n_order - 1] = o0
933933

934+
def _lfilter_core_loop_dispatcher(input_signal_windows: Tensor, a_coeffs_flipped: Tensor, padded_output_waveform: Tensor):
935+
if input_signal_windows.is_cuda or input_signal_windows.is_cpu:
936+
return torch.ops.torchaudio._lfilter_core_loop(input_signal_windows, a_coeffs_flipped, padded_output_waveform)
937+
else:
938+
return _lfilter_core_generic_loop(input_signal_windows, a_coeffs_flipped, padded_output_waveform)
934939

935940
if _IS_TORCHAUDIO_EXT_AVAILABLE:
936-
_lfilter_core_loop = torch.ops.torchaudio._lfilter_core_loop
941+
_lfilter_core_loop = _lfilter_core_loop_dispatcher
937942
else:
938943
_lfilter_core_loop = _lfilter_core_generic_loop
939944

0 commit comments

Comments
 (0)