Support for training on TPU

[saeed-dialpad]

Hi, I'm trying to train a Conformer ASR model with CTC loss on a TPU VM. I set the accelerator: tpu and strategy: auto in the Conformer configuration file. However when I start training I got the following error:

  File "/home/username/.local/lib/python3.10/site-packages/torch/functional.py", line 665, in stft
    return _VF.stft(input, n_fft, hop_length, win_length, window,  # type: ignore[attr-defined]
RuntimeError: Error while lowering: [] aten::as_strided, xla_shape=f32[8,6008,512]{2,1,0}, dynamic_dims: (), size=(8, 6008, 512), stride=(961632, 160, 1), storage_offset=0
Error: torch_xla/csrc/ops/as_strided.cpp:26 : Check failed: storage_offset + slice_size <= input_element_count (24608768 vs. 7693056)

The code runs properly if I set accelerator: cpu. It seems PytorchLightning is not fully compatible with TPU. Can someone help me or point me in the right direction, please?

[titu1994]

We do not support TPU in NeMo codebase. Torch support for xla itself is limited as you see in your logs.

[saeed-dialpad]

Thanks for your answer.