Repository: NVIDIA/NeMo

Issue:
TypeError in RNNTJoint
**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?


Discussion:
This is an interesting find, I'll send a patch tomorrow. Thanks for finding it ! One issue is .item() will constantly perform GPU - CPU transfer so the fix isn't so clean but I suppose it will activate only if it enters the if branch anyway so the cost will be small overall. 

As to your final note, yes, if the batch has extremely different seq lengths then there will be padding and joint cannot avoid computation. There are ways of pre sorting the seq length, split and do fused joint forward, and then remap the IDs to original batch, but it turns out the speedup is quite negligible and memory savings are trivial after doing all that complex work because you usually use a fuser batch size << acoustic batch size, the cost of doing adaptive batch remapping is more significant than just calculating the excess padding on the joint (plus the memory usage is only trivially larger).  (On the order of 50-100 megabytes for a difference of sequence lengths over 200 timesteps)

Resolving patch:
diff --git a/nemo/collections/asr/modules/rnnt.py b/nemo/collections/asr/modules/rnnt.py
--- a/nemo/collections/asr/modules/rnnt.py
+++ b/nemo/collections/asr/modules/rnnt.py
@@ -1289,8 +1289,8 @@ def forward(
                 # Extract the sub batch inputs
                 # sub_enc = encoder_outputs[begin:end, ...]
                 # sub_transcripts = transcripts[begin:end, ...]
-                sub_enc = encoder_outputs.narrow(dim=0, start=begin, length=end - begin)
-                sub_transcripts = transcripts.narrow(dim=0, start=begin, length=end - begin)
+                sub_enc = encoder_outputs.narrow(dim=0, start=begin, length=int(end - begin))
+                sub_transcripts = transcripts.narrow(dim=0, start=begin, length=int(end - begin))
 
                 sub_enc_lens = encoder_lengths[begin:end]
                 sub_transcript_lens = transcript_lengths[begin:end]
@@ -1304,15 +1304,15 @@ def forward(
                     # Reduce encoder length to preserve computation
                     # Encoder: [sub-batch, T, D] -> [sub-batch, T', D]; T' < T
                     if sub_enc.shape[1] != max_sub_enc_length:
-                        sub_enc = sub_enc.narrow(dim=1, start=0, length=max_sub_enc_length)
+                        sub_enc = sub_enc.narrow(dim=1, start=0, length=int(max_sub_enc_length))
 
                     # sub_dec = decoder_outputs[begin:end, ...]  # [sub-batch, U, D]
-                    sub_dec = decoder_outputs.narrow(dim=0, start=begin, length=end - begin)  # [sub-batch, U, D]
+                    sub_dec = decoder_outputs.narrow(dim=0, start=begin, length=int(end - begin))  # [sub-batch, U, D]
 
                     # 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)
+                        sub_dec = sub_dec.narrow(dim=1, start=0, length=int(max_sub_transcript_length + 1))
 
                     # Perform joint => [sub-batch, T', U', V + 1]
                     sub_joint = self.joint(sub_enc, sub_dec)
@@ -1322,7 +1322,7 @@ def forward(
                     # Reduce transcript length to correct alignment
                     # Transcript: [sub-batch, L] -> [sub-batch, L']; L' <= L
                     if sub_transcripts.shape[1] != max_sub_transcript_length:
-                        sub_transcripts = sub_transcripts.narrow(dim=1, start=0, length=max_sub_transcript_length)
+                        sub_transcripts = sub_transcripts.narrow(dim=1, start=0, length=int(max_sub_transcript_length))
 
                     # Compute sub batch loss
                     # preserve loss reduction type
@@ -1758,8 +1758,8 @@ def forward(
             # Extract the sub batch inputs
             # sub_enc = encoder_outputs[begin:end, ...]
             # sub_transcripts = transcripts[begin:end, ...]
-            sub_enc = encoder_outputs.narrow(dim=0, start=begin, length=end - begin)
-            sub_transcripts = transcripts.narrow(dim=0, start=begin, length=end - begin)
+            sub_enc = encoder_outputs.narrow(dim=0, start=begin, length=int(end - begin))
+            sub_transcripts = transcripts.narrow(dim=0, start=begin, length=int(end - begin))
 
             sub_enc_lens = encoder_lengths[begin:end]
             sub_transcript_lens = transcript_lengths[begin:end]
@@ -1773,20 +1773,20 @@ def forward(
                 # Reduce encoder length to preserve computation
                 # Encoder: [sub-batch, T, D] -> [sub-batch, T', D]; T' < T
                 if sub_enc.shape[1] != max_sub_enc_length:
-                    sub_enc = sub_enc.narrow(dim=1, start=0, length=max_sub_enc_length)
+                    sub_enc = sub_enc.narrow(dim=1, start=0, length=int(max_sub_enc_length))
 
                 # sub_dec = decoder_outputs[begin:end, ...]  # [sub-batch, U, D]
-                sub_dec = decoder_outputs.narrow(dim=0, start=begin, length=end - begin)  # [sub-batch, U, D]
+                sub_dec = decoder_outputs.narrow(dim=0, start=begin, length=int(end - begin))  # [sub-batch, U, D]
 
                 # 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)
+                    sub_dec = sub_dec.narrow(dim=1, start=0, length=int(max_sub_transcript_length + 1))
 
                 # Reduce transcript length to correct alignment
                 # Transcript: [sub-batch, L] -> [sub-batch, L']; L' <= L
                 if sub_transcripts.shape[1] != max_sub_transcript_length:
-                    sub_transcripts = sub_transcripts.narrow(dim=1, start=0, length=max_sub_transcript_length)
+                    sub_transcripts = sub_transcripts.narrow(dim=1, start=0, length=int(max_sub_transcript_length))
 
                 # Perform sampled joint => [sub-batch, T', U', {V' < V} + 1}]
                 sub_joint, sub_transcripts_remapped = self.sampled_joint(


Test patch: