Can't setup the environment to run ASR with Transducers tutorial
Hi guys,
I am trying to run the [ASR_with_Transducers](https://github.com/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_Transducers.ipynb) tutorials but it has a version incompatible issue.

At first, I tried to setup the Conda env by following the installation guide and this [instruction](https://github.com/NVIDIA/NeMo/issues/3804#issuecomment-1061569531), the numba already installed from Conda but got this issue:
```
[NeMo W 2022-03-18 22:52:28 nemo_logging:349] /home/tako/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/trainer/configuration_validator.py:291: LightningDeprecationWarning: Base `Callback.on_train_batch_start` hook signature has changed in v1.5. The `dataloader_idx` argument will be removed in v1.7.
      rank_zero_deprecation(
    
[NeMo W 2022-03-18 22:52:28 nemo_logging:349] /home/tako/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/trainer/configuration_validator.py:291: LightningDeprecationWarning: Base `Callback.on_train_batch_end` hook signature has changed in v1.5. The `dataloader_idx` argument will be removed in v1.7.
      rank_zero_deprecation(
    
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2]
[NeMo W 2022-03-18 22:52:33 modelPT:496] The lightning trainer received accelerator: <pytorch_lightning.accelerators.gpu.GPUAccelerator object at 0x7fce77a780d0>. We recommend to use 'ddp' instead.
[NeMo I 2022-03-18 22:52:33 modelPT:587] Optimizer config = Novograd (
    Parameter Group 0
        amsgrad: False
        betas: [0.9, 0.0]
        eps: 1e-08
        grad_averaging: False
        lr: 0.05
        weight_decay: 0.001
    )
[NeMo I 2022-03-18 22:52:33 lr_scheduler:833] Scheduler "<nemo.core.optim.lr_scheduler.CosineAnnealing object at 0x7fce52dbf280>" 
    will be used during training (effective maximum steps = 3000) - 
    Parameters : 
    (warmup_steps: null
    warmup_ratio: null
    min_lr: 1.0e-06
    last_epoch: -1
    max_steps: 3000
    )

  | Name              | Type                              | Params
------------------------------------------------------------------------
0 | preprocessor      | AudioToMelSpectrogramPreprocessor | 0     
1 | encoder           | ConvASREncoder                    | 2.3 M 
2 | decoder           | RNNTDecoder                       | 35.4 K
3 | joint             | RNNTJoint                         | 47.3 K
4 | loss              | RNNTLoss                          | 0     
5 | spec_augmentation | SpectrogramAugmentation           | 0     
6 | wer               | RNNTBPEWER                        | 0     
------------------------------------------------------------------------
2.3 M     Trainable params
0         Non-trainable params
2.3 M     Total params
9.380     Total estimated model params size (MB)
[NeMo I 2022-03-18 22:52:36 rnnt_wer_bpe:231] 
    
[NeMo I 2022-03-18 22:52:36 rnnt_wer_bpe:232] reference :rubout g m e f three nine
[NeMo I 2022-03-18 22:52:36 rnnt_wer_bpe:233] predicted :uuuubuujjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjj
[NeMo I 2022-03-18 22:52:39 rnnt_wer_bpe:231] 
    
[NeMo I 2022-03-18 22:52:39 rnnt_wer_bpe:232] reference :m e l v i n
[NeMo I 2022-03-18 22:52:39 rnnt_wer_bpe:233] predicted :uuuubuuuujjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjj
Epoch 0: 0%
0/60 [00:00<?, ?it/s]
---------------------------------------------------------------------------
CudaAPIError                              Traceback (most recent call last)
File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/cudadrv/driver.py:2188, in Linker.add_ptx(self, ptx, name)
   2187 try:
-> 2188     driver.cuLinkAddData(self.handle, enums.CU_JIT_INPUT_PTX,
   2189                          ptxbuf, len(ptx), namebuf, 0, None, None)
   2190 except CudaAPIError as e:

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/cudadrv/driver.py:303, in Driver._wrap_api_call.<locals>.safe_cuda_api_call(*args)
    302 retcode = libfn(*args)
--> 303 self._check_error(fname, retcode)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/cudadrv/driver.py:343, in Driver._check_error(self, fname, retcode)
    342         raise CudaDriverError("CUDA initialized before forking")
--> 343 raise CudaAPIError(retcode, msg)

CudaAPIError: [222] Call to cuLinkAddData results in UNKNOWN_CUDA_ERROR

During handling of the above exception, another exception occurred:

LinkerError                               Traceback (most recent call last)
Input In [32], in <cell line: 2>()
      1 # Train the model
----> 2 trainer.fit(model)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:740, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, train_dataloader, ckpt_path)
    735     rank_zero_deprecation(
    736         "`trainer.fit(train_dataloader)` is deprecated in v1.4 and will be removed in v1.6."
    737         " Use `trainer.fit(train_dataloaders)` instead. HINT: added 's'"
    738     )
    739     train_dataloaders = train_dataloader
--> 740 self._call_and_handle_interrupt(
    741     self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    742 )

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:685, in Trainer._call_and_handle_interrupt(self, trainer_fn, *args, **kwargs)
    675 r"""
    676 Error handling, intended to be used only for main trainer function entry points (fit, validate, test, predict)
    677 as all errors should funnel through them
   (...)
    682     **kwargs: keyword arguments to be passed to `trainer_fn`
    683 """
    684 try:
--> 685     return trainer_fn(*args, **kwargs)
    686 # TODO: treat KeyboardInterrupt as BaseException (delete the code below) in v1.7
    687 except KeyboardInterrupt as exception:

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:777, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    775 # TODO: ckpt_path only in v1.7
    776 ckpt_path = ckpt_path or self.resume_from_checkpoint
--> 777 self._run(model, ckpt_path=ckpt_path)
    779 assert self.state.stopped
    780 self.training = False

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1199, in Trainer._run(self, model, ckpt_path)
   1196 self.checkpoint_connector.resume_end()
   1198 # dispatch `start_training` or `start_evaluating` or `start_predicting`
-> 1199 self._dispatch()
   1201 # plugin will finalized fitting (e.g. ddp_spawn will load trained model)
   1202 self._post_dispatch()

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1279, in Trainer._dispatch(self)
   1277     self.training_type_plugin.start_predicting(self)
   1278 else:
-> 1279     self.training_type_plugin.start_training(self)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py:202, in TrainingTypePlugin.start_training(self, trainer)
    200 def start_training(self, trainer: "pl.Trainer") -> None:
    201     # double dispatch to initiate the training loop
--> 202     self._results = trainer.run_stage()

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1289, in Trainer.run_stage(self)
   1287 if self.predicting:
   1288     return self._run_predict()
-> 1289 return self._run_train()

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1319, in Trainer._run_train(self)
   1317 self.fit_loop.trainer = self
   1318 with torch.autograd.set_detect_anomaly(self._detect_anomaly):
-> 1319     self.fit_loop.run()

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/base.py:145, in Loop.run(self, *args, **kwargs)
    143 try:
    144     self.on_advance_start(*args, **kwargs)
--> 145     self.advance(*args, **kwargs)
    146     self.on_advance_end()
    147     self.restarting = False

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/fit_loop.py:234, in FitLoop.advance(self)
    231 data_fetcher = self.trainer._data_connector.get_profiled_dataloader(dataloader)
    233 with self.trainer.profiler.profile("run_training_epoch"):
--> 234     self.epoch_loop.run(data_fetcher)
    236     # the global step is manually decreased here due to backwards compatibility with existing loggers
    237     # as they expect that the same step is used when logging epoch end metrics even when the batch loop has
    238     # finished. this means the attribute does not exactly track the number of optimizer steps applied.
    239     # TODO(@carmocca): deprecate and rename so users don't get confused
    240     self.global_step -= 1

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/base.py:145, in Loop.run(self, *args, **kwargs)
    143 try:
    144     self.on_advance_start(*args, **kwargs)
--> 145     self.advance(*args, **kwargs)
    146     self.on_advance_end()
    147     self.restarting = False

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py:193, in TrainingEpochLoop.advance(self, *args, **kwargs)
    190     self.batch_progress.increment_started()
    192     with self.trainer.profiler.profile("run_training_batch"):
--> 193         batch_output = self.batch_loop.run(batch, batch_idx)
    195 self.batch_progress.increment_processed()
    197 # update non-plateau LR schedulers
    198 # update epoch-interval ones only when we are at the end of training epoch

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/base.py:145, in Loop.run(self, *args, **kwargs)
    143 try:
    144     self.on_advance_start(*args, **kwargs)
--> 145     self.advance(*args, **kwargs)
    146     self.on_advance_end()
    147     self.restarting = False

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py:88, in TrainingBatchLoop.advance(self, batch, batch_idx)
     86 if self.trainer.lightning_module.automatic_optimization:
     87     optimizers = _get_active_optimizers(self.trainer.optimizers, self.trainer.optimizer_frequencies, batch_idx)
---> 88     outputs = self.optimizer_loop.run(split_batch, optimizers, batch_idx)
     89 else:
     90     outputs = self.manual_loop.run(split_batch, batch_idx)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/base.py:145, in Loop.run(self, *args, **kwargs)
    143 try:
    144     self.on_advance_start(*args, **kwargs)
--> 145     self.advance(*args, **kwargs)
    146     self.on_advance_end()
    147     self.restarting = False

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:215, in OptimizerLoop.advance(self, batch, *args, **kwargs)
    214 def advance(self, batch: Any, *args: Any, **kwargs: Any) -> None:  # type: ignore[override]
--> 215     result = self._run_optimization(
    216         batch,
    217         self._batch_idx,
    218         self._optimizers[self.optim_progress.optimizer_position],
    219         self.optimizer_idx,
    220     )
    221     if result.loss is not None:
    222         # automatic optimization assumes a loss needs to be returned for extras to be considered as the batch
    223         # would be skipped otherwise
    224         self._outputs[self.optimizer_idx] = result.asdict()

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:266, in OptimizerLoop._run_optimization(self, split_batch, batch_idx, optimizer, opt_idx)
    259         closure()
    261 # ------------------------------
    262 # BACKWARD PASS
    263 # ------------------------------
    264 # gradient update with accumulated gradients
    265 else:
--> 266     self._optimizer_step(optimizer, opt_idx, batch_idx, closure)
    268 result = closure.consume_result()
    270 if result.loss is not None:
    271     # if no result, user decided to skip optimization
    272     # otherwise update running loss + reset accumulated loss
    273     # TODO: find proper way to handle updating running loss

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:378, in OptimizerLoop._optimizer_step(self, optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
    375 self.optim_progress.optimizer.step.increment_ready()
    377 # model hook
--> 378 lightning_module.optimizer_step(
    379     self.trainer.current_epoch,
    380     batch_idx,
    381     optimizer,
    382     opt_idx,
    383     train_step_and_backward_closure,
    384     on_tpu=(self.trainer._device_type == DeviceType.TPU and _TPU_AVAILABLE),
    385     using_native_amp=(self.trainer.amp_backend is not None and self.trainer.amp_backend == AMPType.NATIVE),
    386     using_lbfgs=is_lbfgs,
    387 )
    389 self.optim_progress.optimizer.step.increment_completed()

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/core/lightning.py:1652, in LightningModule.optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure, on_tpu, using_native_amp, using_lbfgs)
   1571 def optimizer_step(
   1572     self,
   1573     epoch: int,
   (...)
   1580     using_lbfgs: bool = False,
   1581 ) -> None:
   1582     r"""
   1583     Override this method to adjust the default way the
   1584     :class:`~pytorch_lightning.trainer.trainer.Trainer` calls each optimizer.
   (...)
   1650 
   1651     """
-> 1652     optimizer.step(closure=optimizer_closure)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/core/optimizer.py:164, in LightningOptimizer.step(self, closure, **kwargs)
    162 assert trainer is not None
    163 with trainer.profiler.profile(profiler_action):
--> 164     trainer.accelerator.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/accelerators/accelerator.py:339, in Accelerator.optimizer_step(self, optimizer, opt_idx, closure, model, **kwargs)
    329 """performs the actual optimizer step.
    330 
    331 Args:
   (...)
    336     **kwargs: Any extra arguments to ``optimizer.step``
    337 """
    338 model = model or self.lightning_module
--> 339 self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py:163, in PrecisionPlugin.optimizer_step(self, model, optimizer, optimizer_idx, closure, **kwargs)
    161 if isinstance(model, pl.LightningModule):
    162     closure = partial(self._wrap_closure, model, optimizer, optimizer_idx, closure)
--> 163 optimizer.step(closure=closure, **kwargs)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/torch/optim/lr_scheduler.py:65, in _LRScheduler.__init__.<locals>.with_counter.<locals>.wrapper(*args, **kwargs)
     63 instance._step_count += 1
     64 wrapped = func.__get__(instance, cls)
---> 65 return wrapped(*args, **kwargs)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/torch/optim/optimizer.py:88, in Optimizer._hook_for_profile.<locals>.profile_hook_step.<locals>.wrapper(*args, **kwargs)
     86 profile_name = "Optimizer.step#{}.step".format(obj.__class__.__name__)
     87 with torch.autograd.profiler.record_function(profile_name):
---> 88     return func(*args, **kwargs)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/core/optim/novograd.py:83, in Novograd.step(self, closure)
     81 loss = None
     82 if closure is not None:
---> 83     loss = closure()
     85 for group in self.param_groups:
     86     for p in group["params"]:

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py:148, in PrecisionPlugin._wrap_closure(self, model, optimizer, optimizer_idx, closure)
    135 def _wrap_closure(
    136     self,
    137     model: "pl.LightningModule",
   (...)
    140     closure: Callable[[], Any],
    141 ) -> Any:
    142     """This double-closure allows makes sure the ``closure`` is executed before the
    143     ``on_before_optimizer_step`` hook is called.
    144 
    145     The closure (generally) runs ``backward`` so this allows inspecting gradients in this hook. This structure is
    146     consistent with the ``PrecisionPlugin`` subclasses that cannot pass ``optimizer.step(closure)`` directly.
    147     """
--> 148     closure_result = closure()
    149     self._after_closure(model, optimizer, optimizer_idx)
    150     return closure_result

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:160, in Closure.__call__(self, *args, **kwargs)
    159 def __call__(self, *args: Any, **kwargs: Any) -> Optional[Tensor]:
--> 160     self._result = self.closure(*args, **kwargs)
    161     return self._result.loss

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:142, in Closure.closure(self, *args, **kwargs)
    140 def closure(self, *args: Any, **kwargs: Any) -> ClosureResult:
    141     with self._profiler.profile("training_step_and_backward"):
--> 142         step_output = self._step_fn()
    144         if step_output.closure_loss is None:
    145             self.warning_cache.warn(
    146                 "`training_step` returned `None`. If this was on purpose, ignore this warning..."
    147             )

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py:435, in OptimizerLoop._training_step(self, split_batch, batch_idx, opt_idx)
    433 lightning_module._current_fx_name = "training_step"
    434 with self.trainer.profiler.profile("training_step"):
--> 435     training_step_output = self.trainer.accelerator.training_step(step_kwargs)
    436     self.trainer.training_type_plugin.post_training_step()
    438 del step_kwargs

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/accelerators/accelerator.py:219, in Accelerator.training_step(self, step_kwargs)
    214 """The actual training step.
    215 
    216 See :meth:`~pytorch_lightning.core.lightning.LightningModule.training_step` for more details
    217 """
    218 with self.precision_plugin.train_step_context():
--> 219     return self.training_type_plugin.training_step(*step_kwargs.values())

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py:213, in TrainingTypePlugin.training_step(self, *args, **kwargs)
    212 def training_step(self, *args, **kwargs):
--> 213     return self.model.training_step(*args, **kwargs)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/utils/model_utils.py:355, in wrap_training_step(wrapped, instance, args, kwargs)
    353 @wrapt.decorator
    354 def wrap_training_step(wrapped, instance: 'pl.LightningModule', args, kwargs):
--> 355     output_dict = wrapped(*args, **kwargs)
    357     if isinstance(output_dict, dict) and output_dict is not None and 'log' in output_dict:
    358         log_dict = output_dict.pop('log')

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/models/rnnt_models.py:711, in EncDecRNNTModel.training_step(self, batch, batch_nb)
    708     compute_wer = False
    710 # Fused joint step
--> 711 loss_value, wer, _, _ = self.joint(
    712     encoder_outputs=encoded,
    713     decoder_outputs=decoder,
    714     encoder_lengths=encoded_len,
    715     transcripts=transcript,
    716     transcript_lengths=transcript_len,
    717     compute_wer=compute_wer,
    718 )
    720 tensorboard_logs = {'train_loss': loss_value, 'learning_rate': self._optimizer.param_groups[0]['lr']}
    722 if compute_wer:

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/torch/nn/modules/module.py:1110, in Module._call_impl(self, *input, **kwargs)
   1106 # If we don't have any hooks, we want to skip the rest of the logic in
   1107 # this function, and just call forward.
   1108 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110     return forward_call(*input, **kwargs)
   1111 # Do not call functions when jit is used
   1112 full_backward_hooks, non_full_backward_hooks = [], []

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/core/classes/common.py:822, in typecheck.__call__(self, wrapped, instance, args, kwargs)
    819 instance._validate_input_types(input_types=input_types, ignore_collections=self.ignore_collections, **kwargs)
    821 # Call the method - this can be forward, or any other callable method
--> 822 outputs = wrapped(*args, **kwargs)
    824 instance._attach_and_validate_output_types(
    825     output_types=output_types, ignore_collections=self.ignore_collections, out_objects=outputs
    826 )
    828 return outputs

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/modules/rnnt.py:885, in RNNTJoint.forward(self, encoder_outputs, decoder_outputs, encoder_lengths, transcripts, transcript_lengths, compute_wer)
    882 self.loss.reduction = None
    884 # compute and preserve loss
--> 885 loss_batch = self.loss(
    886     log_probs=sub_joint,
    887     targets=sub_transcripts,
    888     input_lengths=sub_enc_lens,
    889     target_lengths=sub_transcript_lens,
    890 )
    891 losses.append(loss_batch)
    893 # reset loss reduction type

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/torch/nn/modules/module.py:1110, in Module._call_impl(self, *input, **kwargs)
   1106 # If we don't have any hooks, we want to skip the rest of the logic in
   1107 # this function, and just call forward.
   1108 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110     return forward_call(*input, **kwargs)
   1111 # Do not call functions when jit is used
   1112 full_backward_hooks, non_full_backward_hooks = [], []

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/core/classes/common.py:822, in typecheck.__call__(self, wrapped, instance, args, kwargs)
    819 instance._validate_input_types(input_types=input_types, ignore_collections=self.ignore_collections, **kwargs)
    821 # Call the method - this can be forward, or any other callable method
--> 822 outputs = wrapped(*args, **kwargs)
    824 instance._attach_and_validate_output_types(
    825     output_types=output_types, ignore_collections=self.ignore_collections, out_objects=outputs
    826 )
    828 return outputs

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/losses/rnnt.py:288, in RNNTLoss.forward(self, log_probs, targets, input_lengths, target_lengths)
    285     self._loss.reduction = self.reduction
    287 # Compute RNNT loss
--> 288 loss = self._loss(acts=log_probs, labels=targets, act_lens=input_lengths, label_lens=target_lengths)
    290 # Loss reduction can be dynamic, so reset it after call
    291 if self.reduction != 'mean_batch':

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/torch/nn/modules/module.py:1110, in Module._call_impl(self, *input, **kwargs)
   1106 # If we don't have any hooks, we want to skip the rest of the logic in
   1107 # this function, and just call forward.
   1108 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1109         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1110     return forward_call(*input, **kwargs)
   1111 # Do not call functions when jit is used
   1112 full_backward_hooks, non_full_backward_hooks = [], []

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/rnnt_pytorch.py:166, in RNNTLossNumba.forward(self, acts, labels, act_lens, label_lens)
    162     # NOTE manually done log_softmax for CPU version,
    163     # log_softmax is computed within GPU version.
    164     acts = torch.nn.functional.log_softmax(acts, -1)
--> 166 return self.loss(
    167     acts, labels, act_lens, label_lens, self.blank, self.reduction, self.fastemit_lambda, self.clamp
    168 )

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/rnnt_pytorch.py:62, in _RNNTNumba.forward(ctx, acts, labels, act_lens, label_lens, blank, reduction, fastemit_lambda, clamp)
     59 minibatch_size = acts.size(0)
     60 costs = torch.zeros(minibatch_size, device=acts.device, dtype=acts.dtype)
---> 62 loss_func(
     63     acts,
     64     labels=labels,
     65     input_lengths=act_lens,
     66     label_lengths=label_lens,
     67     costs=costs,
     68     grads=grads,
     69     blank_label=blank,
     70     fastemit_lambda=fastemit_lambda,
     71     clamp=clamp,
     72     num_threads=0,
     73 )
     75 if reduction in ['sum', 'mean']:
     76     costs = costs.sum().unsqueeze_(-1)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/rnnt.py:223, in rnnt_loss_gpu(acts, labels, input_lengths, label_lengths, costs, grads, blank_label, fastemit_lambda, clamp, num_threads)
    219 else:
    220     ### FLATTEN GRAD TENSOR ###
    221     grads, grads_shape = rnnt_helper.flatten_tensor(grads)
--> 223     status = wrapper.cost_and_grad(
    224         acts=acts.data,
    225         grads=grads.data,
    226         costs=costs.data,
    227         pad_labels=labels.data,
    228         label_lengths=label_lengths.data,
    229         input_lengths=input_lengths.data,
    230     )
    232     if status != global_constants.RNNTStatus.RNNT_STATUS_SUCCESS:
    233         raise RuntimeError("Could not calculate forward scores")

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt.py:249, in GPURNNT.cost_and_grad(self, acts, grads, costs, pad_labels, label_lengths, input_lengths)
    239 if (
    240     acts is None
    241     or grads is None
   (...)
    245     or input_lengths is None
    246 ):
    247     return global_constants.RNNTStatus.RNNT_STATUS_INVALID_VALUE
--> 249 return self.compute_cost_and_score(acts, grads, costs, pad_labels, label_lengths, input_lengths)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt.py:158, in GPURNNT.compute_cost_and_score(self, acts, grads, costs, labels, label_lengths, input_lengths)
    155 used_offset, (denom, alphas, betas, llForward, llBackward) = self._prepare_workspace()
    157 ######## START EXECUTION ########
--> 158 self.log_softmax(acts, denom)
    160 # Compute alphas
    161 gpu_rnnt_kernel.compute_alphas_kernel[self.minibatch_, self.maxU_, self.stream_, 0](
    162     acts,
    163     denom,
   (...)
    173     self.blank_,
    174 )

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt.py:104, in GPURNNT.log_softmax(self, acts, denom)
     91 """
     92 Computes the log softmax denominator of the input activation tensor
     93 and stores the result in denom.
   (...)
    101     This kernel inplace updates the `denom` tensor
    102 """
    103 # // trans_acts + pred_acts -> log_softmax denominator
--> 104 reduce.reduce_max(
    105     acts,
    106     denom,
    107     rows=self.alphabet_size_,
    108     cols=self.minibatch_ * self.maxT_ * self.maxU_,
    109     minus=False,
    110     stream=self.stream_,
    111 )
    113 reduce.reduce_exp(
    114     acts,
    115     denom,
   (...)
    119     stream=self.stream_,
    120 )

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/reduce.py:353, in reduce_max(acts, denom, rows, cols, minus, stream)
    332 def reduce_max(acts: torch.Tensor, denom, rows: int, cols: int, minus: bool, stream):
    333     """
    334     Helper method to call the Warp Reduction Kernel to perform `max` reduction.
    335 
   (...)
    351         stream: CUDA Stream.
    352     """
--> 353     return ReduceHelper(
    354         I_opid=I_Op.IDENTITY.value,
    355         R_opid=R_Op.MAXIMUM.value,
    356         acts=acts,
    357         output=denom,
    358         num_rows=rows,
    359         num_cols=cols,
    360         minus=minus,
    361         stream=stream,
    362     )

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/reduce.py:294, in ReduceHelper(I_opid, R_opid, acts, output, num_rows, num_cols, minus, stream)
    292     grid_size = num_cols
    293     # call kernel
--> 294     _reduce_rows[grid_size, CTA_REDUCE_SIZE, stream, 0](I_opid, R_opid, acts, output, num_rows)
    296 return True

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/compiler.py:868, in _KernelConfiguration.__call__(self, *args)
    867 def __call__(self, *args):
--> 868     return self.dispatcher.call(args, self.griddim, self.blockdim,
    869                                 self.stream, self.sharedmem)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/compiler.py:1001, in Dispatcher.call(self, args, griddim, blockdim, stream, sharedmem)
    999     kernel = next(iter(self.overloads.values()))
   1000 else:
-> 1001     kernel = _dispatcher.Dispatcher._cuda_call(self, *args)
   1003 kernel.launch(args, griddim, blockdim, stream, sharedmem)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/compiler.py:1009, in Dispatcher._compile_for_args(self, *args, **kws)
   1007 assert not kws
   1008 argtypes = [self.typeof_pyval(a) for a in args]
-> 1009 return self.compile(tuple(argtypes))

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/compiler.py:1115, in Dispatcher.compile(self, sig)
   1112     self._insert(c_sig, kernel, cuda=True)
   1113     self.overloads[argtypes] = kernel
-> 1115     kernel.bind()
   1116     self.sigs.append(sig)
   1117 return kernel

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/compiler.py:626, in _Kernel.bind(self)
    622 def bind(self):
    623     """
    624     Force binding to current CUDA context
    625     """
--> 626     self._codelibrary.get_cufunc()

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/codegen.py:201, in CUDACodeLibrary.get_cufunc(self)
    198 if cufunc:
    199     return cufunc
--> 201 cubin = self.get_cubin(cc=device.compute_capability)
    202 module = ctx.create_module_image(cubin)
    204 # Load

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/codegen.py:174, in CUDACodeLibrary.get_cubin(self, cc)
    172 ptxes = self._get_ptxes(cc=cc)
    173 for ptx in ptxes:
--> 174     linker.add_ptx(ptx.encode())
    175 for path in self._linking_files:
    176     linker.add_file_guess_ext(path)

File ~/anaconda3/envs/nemo_stt/lib/python3.9/site-packages/numba/cuda/cudadrv/driver.py:2191, in Linker.add_ptx(self, ptx, name)
   2188     driver.cuLinkAddData(self.handle, enums.CU_JIT_INPUT_PTX,
   2189                          ptxbuf, len(ptx), namebuf, 0, None, None)
   2190 except CudaAPIError as e:
-> 2191     raise LinkerError("%s\n%s" % (e, self.error_log))

LinkerError: [222] Call to cuLinkAddData results in UNKNOWN_CUDA_ERROR
ptxas application ptx input, line 9; fatal   : Unsupported .version 7.3; current version is '7.2'
```


GPU infos:
```
Fri Mar 18 23:12:27 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.80       Driver Version: 460.80       CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  GeForce RTX 3090    Off  | 00000000:3B:00.0 Off |                  N/A |
| 54%   25C    P8    25W / 350W |      0MiB / 24268MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  GeForce RTX 3090    Off  | 00000000:5E:00.0 Off |                  N/A |
| 58%   25C    P8    23W / 350W |      0MiB / 24268MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  GeForce RTX 3090    Off  | 00000000:AF:00.0 Off |                  N/A |
|  0%   29C    P8    27W / 350W |      0MiB / 24268MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

```

nvcc --version:
'''
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Sun_Mar_21_19:15:46_PDT_2021
Cuda compilation tools, release 11.3, V11.3.58
Build cuda_11.3.r11.3/compiler.29745058_0
'''


