TypeError in RNNTJoint
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**
Hi. In my scenario `RNNTJoint.forward` produces TypeError and error message is 
```
narrow() received an invalid combination of arguments - got (length=Tensor, start=int, dim=int, ), but expected one of:
 * (int dim, Tensor start, int length)
      didn't match because some of the arguments have invalid types: (dim=int, start=int, length=Tensor, )
 * (int dim, int start, int length)
      didn't match because some of the arguments have invalid types: (dim=int, start=int, length=Tensor )
```

**Steps/Code to reproduce bug**
I faced this issue with nemo-toolkit==1.12 but it looks like problem still exists in the version 1.16. 
The nature of the error itself can be shown using the following code:
```python
import torch
a = torch.Tensor([0,1,1])
max_a = a.max()
a.narrow(dim=0, start=0, length = max_a)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-1-1104b945c0b8> in <cell line: 4>()
      2 a = torch.Tensor([0,1,1])
      3 max_a = a.max()
----> 4 a.narrow(dim=0, start=0, length = max_a)

TypeError: narrow() received an invalid combination of arguments - got (length=Tensor, start=int, dim=int, ), but expected one of:
 * (int dim, Tensor start, int length)
      didn't match because some of the arguments have invalid types: (dim=int, !start=int!, !length=Tensor!, )
 * (int dim, int start, int length)
      didn't match because some of the arguments have invalid types: (dim=int, start=int, !length=Tensor!, )
```

And there is the similar code in the [nemo 1.12](https://github.com/NVIDIA/NeMo/blob/dd9a30f55b7ec308cc3b55751e8ae5aa40f36a62/nemo/collections/asr/modules/rnnt.py#L1311
) (lines 1297 and 1311) as well as in the [nemo 1.16](https://github.com/NVIDIA/NeMo/blob/163111870e78a2a7224953a89b6af5ac71cb8ed9/nemo/collections/asr/modules/rnnt.py#L1315) (lines 1301 and 1315)

In order to reproduce this behavior you have to make the condition `sub_dec.shape[1] != max_sub_transcript_length + 1`  true. In my scenario, this is true because I use several hacks around my own custom dataset and dataloader that applies padding to the data if "time" dimension of tensor  is not a multiple of 128. But I'm quite sure there is a cleaner way to fulfill the condition, and it looks like the code will throw this error every time the condition is met.

Replacement `max_sub_transcript_length = sub_transcript_lens.max()` -> `max_sub_transcript_length = int(sub_transcript_lens.max().item())` applied [here](https://github.com/NVIDIA/NeMo/blob/dd9a30f55b7ec308cc3b55751e8ae5aa40f36a62/nemo/collections/asr/modules/rnnt.py#L1297) fixes the issue for me.

**Expected behavior**

The line `sub_dec = sub_dec.narrow(dim=1, start=0, length=max_sub_transcript_length + 1)` can be executed without rising TypeError

**Environment overview (please complete the following information)**

 - Environment location: Docker, image is based on cuda:11.3.1-cudnn8-devel-ubuntu20.04
 - Method of NeMo install: `pip install nemo-toolkit['all']==1.12.0`

**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- Ubuntu20.04
- PyTorch version 1.10.1 (but the cause of the problem persists in 1.13.1 )
- Python version 3.9.7

**Additional context**
I have an additional question about the code I mentioned in this issue. There is a couple of comments that says:
```
# Reduce decoder length to preserve computation
# Decoder: [sub-batch, U, D] -> [sub-batch, U', D]; U' < U
    if sub_dec.shape[1] != max_sub_transcript_length + 1:
        sub_dec = sub_dec.narrow(dim=1, start=0, length=max_sub_transcript_length + 1)
```
I understand that we can remove the padding here in order to reduce memory consumption and computation time of the joint layer but we (probably) still have a lot of padding in the decoder output if our batch contains transcriptions of different lengths. As afs as I know current implementation of RNNT will process this padding in the joint layer and there is no way to reduce related computation and memory costs. Is it correct?