transforms.MFCC results in NaN values on Jetson Orin Nano

[frmser]

🐛 Describe the bug

On a Jetson Orin Nano, applying the transforms.MFCC transform with certain parameters to samples from the datasets.SPEECHCOMMANDS dataset, results in some values in the resulting tensors to be NaN. I have tested this script on a few machines with torch and torchaudio 2.1.0 and this does not happen on any of them. As they are all x86, this leads me to believe the bug could be related to ARM processors. I followed the instructions when building torchaudio and did not do anything unusual there. Also, not specifying parameters does not result in the bug but I am used to use these parameters in experiments as they perform better than the default ones in my experience.

import torch
import torchaudio


speech_commands = torchaudio.datasets.SPEECHCOMMANDS(root="data", download=True)

for i in range(len(speech_commands)):
    waveform, sample_rate, _, _, _ = speech_commands.__getitem__(i)

    assert not torch.isnan(waveform).any(), "NaN in waveform"

    mfcc = torchaudio.transforms.MFCC(
                    sample_rate=sample_rate,
                    n_mfcc=40,
                    melkwargs={
                        "n_fft": 480,
                        "n_mels": 40,
                        "win_length": 480,
                        "hop_length": 160,
                        "center": True,
                    },
                )(waveform)

    assert not torch.isnan(mfcc).any(), f"NaN in mfcc {i}"

Traceback (most recent call last):
  File "mfcc_test.py", line 24, in <module>
    assert not torch.isnan(mfcc).any(), f"NaN in mfcc {i}"
AssertionError: NaN in mfcc 0

Versions

Collecting environment information...
PyTorch version: 2.1.0a0+41361538.nv23.06
Is debug build: False
CUDA used to build PyTorch: 11.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (aarch64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.30.0
Libc version: glibc-2.31

Python version: 3.8.10 (default, Mar 25 2024, 10:42:49) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.10.120-tegra-aarch64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/aarch64-linux-gnu/libcudnn.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_adv_infer.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_adv_train.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_cnn_infer.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_cnn_train.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_ops_infer.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_ops_train.so.8.6.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False

CPU:
Architecture: aarch64
CPU op-mode(s): 32-bit, 64-bit
Byte Order: Little Endian
CPU(s): 6
On-line CPU(s) list: 0-5
Thread(s) per core: 1
Core(s) per socket: 6
Socket(s): 1
Vendor ID: ARM
Model: 1
Model name: ARMv8 Processor rev 1 (v8l)
Stepping: r0p1
CPU max MHz: 1510,4000
CPU min MHz: 115,2000
BogoMIPS: 62.50
L1d cache: 384 KiB
L1i cache: 384 KiB
L2 cache: 1,5 MiB
L3 cache: 6 MiB
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, but not BHB
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Not affected
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm lrcpc dcpop asimddp uscat ilrcpc flagm

Versions of relevant libraries:
[pip3] mypy==1.10.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.4
[pip3] torch==2.1.0a0+41361538.nv23.6
[pip3] torchaudio==2.1.0+6ea1133
[pip3] torchvision==0.16.1+fdea156
[conda] Could not collect