# Buggy code
def _feature_window_function(
    window_type: str, window_size: int, blackman_coeff: float
) -> Tensor:
    r"""Returns a window function with the given type and size"""
    if window_type == HANNING:
        return torch.hann_window(window_size, periodic=False)
    elif window_type == HAMMING:
        return torch.hamming_window(window_size, periodic=False, alpha=0.54, beta=0.46)
    elif window_type == POVEY:
        # like hanning but goes to zero at edges
        return torch.hann_window(window_size, periodic=False).pow(0.85)
    elif window_type == RECTANGULAR:
        return torch.ones(window_size)
    elif window_type == BLACKMAN:
        a = 2 * math.pi / (window_size - 1)
        window_function = torch.arange(window_size)
        # can't use torch.blackman_window as they use different coefficients
        return (
            blackman_coeff
            - 0.5 * torch.cos(a * window_function)
            + (0.5 - blackman_coeff) * torch.cos(2 * a * window_function)
        )
    else:
        raise Exception("Invalid window type " + window_type)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "./fbank_test.py", line 27, in <module>
main()
File "./fbank_test.py", line 23, in main
fbank = torchaudio.compliance.kaldi.fbank(waveform)
File "/xxxxxx/py35/lib/python3.5/site-packages/torchaudio/compliance/kaldi.py",
line 554, in fbank
snip_edges, raw_energy, energy_floor, dither, remove_dc_offset, preemphasis_coefficient)
File "/xxxxx/py35/lib/python3.5/site-packages/torchaudio/compliance/kaldi.py",
line 180, in _get_window
signal_log_energy = _get_log_energy(strided_input, EPSILON, energy_floor)  # size (m)
File "/xxxxx/py35/lib/python3.5/site-packages/torchaudio/compliance/kaldi.py",
line 113, in _get_log_energy
log_energy = torch.max(strided_input.pow(2).sum(1), epsilon).log()  # size (m)
RuntimeError: iter.device(arg).is_cuda() INTERNAL ASSERT FAILED at /pytorch/aten/src/ATen/native/cuda/Loops.cuh:56,
please report a bug to PyTorch.

# Fixed code
def _feature_window_function(
    window_type: str,
    window_size: int,
    blackman_coeff: float,
    device: torch.device,
    dtype: int,
) -> Tensor:
    r"""Returns a window function with the given type and size"""
    if window_type == HANNING:
        return torch.hann_window(
            window_size, periodic=False, device=device, dtype=dtype
        )
    elif window_type == HAMMING:
        return torch.hamming_window(
            window_size,
            periodic=False,
            alpha=0.54,
            beta=0.46,
            device=device,
            dtype=dtype,
        )
    elif window_type == POVEY:
        # like hanning but goes to zero at edges
        return torch.hann_window(
            window_size, periodic=False, device=device, dtype=dtype
        ).pow(0.85)
    elif window_type == RECTANGULAR:
        return torch.ones(window_size, device=device, dtype=dtype)
    elif window_type == BLACKMAN:
        a = 2 * math.pi / (window_size - 1)
        window_function = torch.arange(window_size, device=device, dtype=dtype)
        # can't use torch.blackman_window as they use different coefficients
        return (
            blackman_coeff
            - 0.5 * torch.cos(a * window_function)
            + (0.5 - blackman_coeff) * torch.cos(2 * a * window_function)
        ).to(device=device, dtype=dtype)
    else:
        raise Exception("Invalid window type " + window_type)