Numba error while fine-tuning an ASR model
Hello,

I am fine-tuning an ASR model (the English Conformer Transducer) and encountering the following error:

```log
2023-02-03 14:46:28 | numba.cuda.cudadrv.driver.CudaAPIError: [1] Call to cuLaunchKernel results in CUDA_ERROR_INVALID_VALUE |  
-- | -- | --
  |   | 2023-02-03 14:46:28 | raise CudaAPIError(retcode, msg) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py", line 335, in _check_error |  
  |   | 2023-02-03 14:46:28 | self._check_error(fname, retcode) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py", line 300, in safe_cuda_api_call |  
  |   | 2023-02-03 14:46:28 | driver.cuLaunchKernel(cufunc_handle, |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py", line 2094, in launch_kernel |  
  |   | 2023-02-03 14:46:28 | driver.launch_kernel(cufunc.handle, |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/numba/cuda/compiler.py", line 693, in launch |  
  |   | 2023-02-03 14:46:28 | kernel.launch(args, griddim, blockdim, stream, sharedmem) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/numba/cuda/compiler.py", line 966, in call |  
  |   | 2023-02-03 14:46:28 | return self.dispatcher.call(args, self.griddim, self.blockdim, |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/numba/cuda/compiler.py", line 821, in __call__ |  
  |   | 2023-02-03 14:46:28 | gpu_rnnt_kernel.compute_alphas_kernel[self.minibatch_, self.maxU_, self.stream_, 0]( |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt.py", line 161, in compute_cost_and_score |  
  |   | 2023-02-03 14:46:28 | return self.compute_cost_and_score(acts, grads, costs, pad_labels, label_lengths, input_lengths) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt.py", line 249, in cost_and_grad |  
  |   | 2023-02-03 14:46:28 | status = wrapper.cost_and_grad( |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py", line 223, in rnnt_loss_gpu |  
  |   | 2023-02-03 14:46:28 | loss_func( |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/rnnt_pytorch.py", line 62, in forward |  
  |   | 2023-02-03 14:46:28 | return self.loss( |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/rnnt_pytorch.py", line 166, in forward |  
  |   | 2023-02-03 14:46:28 | return forward_call(*input, **kwargs) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1111, in _call_impl |  
  |   | 2023-02-03 14:46:28 | loss = self._loss(acts=log_probs, labels=targets, act_lens=input_lengths, label_lens=target_lengths) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/losses/rnnt.py", line 288, in forward |  
  |   | 2023-02-03 14:46:28 | loss_batch = self.loss( |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/modules/rnnt.py", line 901, in forward |  
  |   | 2023-02-03 14:46:28 | outputs = wrapped(*args, **kwargs) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/nemo/core/classes/common.py", line 963, in __call__ |  
  |   | 2023-02-03 14:46:28 | loss_value, wer, _, _ = self.joint( |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/nemo/collections/asr/models/rnnt_models.py", line 710, in training_step |  
  |   | 2023-02-03 14:46:28 | output_dict = wrapped(*args, **kwargs) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/nemo/utils/model_utils.py", line 364, in wrap_training_step |  
  |   | 2023-02-03 14:46:28 | asr_loss = self.asr.training_step(batch, batch_idx)["loss"] |  
  |   | 2023-02-03 14:46:28 | File "/netscratch/franzreb/accent_info/src/model.py", line 121, in training_step |  
  |   | 2023-02-03 14:46:28 | output = self.module.training_step(*inputs, **kwargs) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 82, in forward |  
  |   | 2023-02-03 14:46:28 | output = self.module(*inputs[0], **kwargs[0]) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 992, in forward |  
  |   | 2023-02-03 14:46:28 | return forward_call(*input, **kwargs) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1111, in _call_impl |  
  |   | 2023-02-03 14:46:28 | return self.model(*args, **kwargs) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 349, in training_step |  
  |   | 2023-02-03 14:46:28 | output = fn(*args, **kwargs) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1765, in _call_strategy_hook |  
  |   | 2023-02-03 14:46:28 | training_step_output = self.trainer._call_strategy_hook("training_step", *step_kwargs.values()) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 427, in _training_step |  
  |   | 2023-02-03 14:46:28 | step_output = self._step_fn() |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 134, in closure |  
  |   | 2023-02-03 14:46:28 | self._result = self.closure(*args, **kwargs) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 148, in __call__ |  
  |   | 2023-02-03 14:46:28 | closure() |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 249, in _run_optimization |  
  |   | 2023-02-03 14:46:28 | result = self._run_optimization( |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 203, in advance |  
  |   | 2023-02-03 14:46:28 | outputs = self.optimizer_loop.run(split_batch, optimizers, batch_idx) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 88, in advance |  
  |   | 2023-02-03 14:46:28 | batch_output = self.batch_loop.run(batch, batch_idx) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 208, in advance |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run |  
  |   | 2023-02-03 14:46:28 | self._outputs = self.epoch_loop.run(self._data_fetcher) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py", line 269, in advance |  
  |   | 2023-02-03 14:46:28 | self.advance(*args, **kwargs) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run |  
  |   | 2023-02-03 14:46:28 | self.fit_loop.run() |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1353, in _run_train |  
  |   | 2023-02-03 14:46:28 | return self._run_train() |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1323, in _run_stage |  
  |   | 2023-02-03 14:46:28 | results = self._run_stage() |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1236, in _run |  
  |   | 2023-02-03 14:46:28 | results = self._run(model, ckpt_path=self.ckpt_path) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 811, in _fit_impl |  
  |   | 2023-02-03 14:46:28 | return trainer_fn(*args, **kwargs) |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 723, in _call_and_handle_interrupt |  
  |   | 2023-02-03 14:46:28 | self._call_and_handle_interrupt( |  
  |   | 2023-02-03 14:46:28 | File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 770, in fit |  
  |   | 2023-02-03 14:46:28 | trainer.fit(model, train_dataloaders=train_dl) |  
  |   | 2023-02-03 14:46:28 | File "run.py", line 20, in main |  
  |   | 2023-02-03 14:46:28 | main(args) |  
  |   | 2023-02-03 14:46:28 | File "run.py", line 57, in <module> |  
  |   | 2023-02-03 14:46:28 | Traceback (most recent call last):
```

The error always occurs at a similar training step. It is a Numba Cuda API error that happens when calling the ASR model's loss function. Do you know how to address this issue? If not, how can I debug it? Thank you very much for your help.