'CUDA error: an illegal memory access was encountered' when using forced_align on cuda device > 0

[Approximetal]

🐛 Describe the bug

It seems torch.ops.torchaudio.forced_align function does not support cuda:n when n>0, if I set device='cuda:0', it works well, but if I set device='cuda:1', I got errors like this:

"""
Traceback (most recent call last):
  File "/mnt/miniconda/envs/cv/lib/python3.10/concurrent/futures/process.py", line 246, in _process_worker
    r = call_item.fn(*call_item.args, **call_item.kwargs)
  File "/mnt/code/speechTranslate/process_chain/gen_tts.py", line 248, in process_batch
    main_process(data["ref_audio"], data["ref_text"], data["gen_text"], data["save_path"], f5_models, mms_models, vc_models, data["speed"], device)
  File "/mnt/code/speechTranslate/process_chain/gen_tts.py", line 117, in main_process
    align_res = mms_align((audio,24000), gen_text, mms_model, tokenizer, mms_aligner, token_dict, text_normalizer, device=device)
  File "/mnt/code/speechTranslate/process_chain/gen_align.py", line 100, in process_one
    token_spans = mms_aligner(emission[0], tokens)
  File "/mnt/miniconda/envs/cv/lib/python3.10/site-packages/torchaudio/pipelines/_wav2vec2/aligner.py", line 85, in __call__
    aligned_tokens, scores = _align_emission_and_tokens(emission, _flatten(tokens), self.blank)
  File "/mnt/miniconda/envs/cv/lib/python3.10/site-packages/torchaudio/pipelines/_wav2vec2/aligner.py", line 40, in _align_emission_and_tokens
    aligned_tokens, scores = F.forced_align(emission, targets, blank=blank)
  File "/mnt/miniconda/envs/cv/lib/python3.10/site-packages/torchaudio/functional/_alignment.py", line 72, in forced_align
    paths, scores = torch.ops.torchaudio.forced_align(log_probs, targets, input_lengths, target_lengths, blank)
  File "/mnt/miniconda/envs/cv/lib/python3.10/site-packages/torch/_ops.py", line 854, in __call__
    return self_._op(*args, **(kwargs or {}))
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/mnt/code/speechTranslate/process_chain/gen_tts.py", line 368, in <module>
    parallel_process(filenames, num_processes, args)
  File "/mnt/code/speechTranslate/process_chain/gen_tts.py", line 271, in parallel_process
    task.result()
  File "/mnt/miniconda/envs/cv/lib/python3.10/concurrent/futures/_base.py", line 458, in result
    return self.__get_result()
  File "/mnt/miniconda/envs/cv/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result
    raise self._exception
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Versions

multiple version of torchaudio shows this error.

I found the same issue in pytorch/pytorch#148438, and similar solutions in
#3425 and https://discuss.pytorch.org/t/c-cuda-extension-with-multiple-gpus/91241/5?u=yoyololicon
but I'm not sure where should I change to fix this function torch.ops.torchaudio.forced_align