The traced Encoder of LEDForConditionalGeneration does not allow dynamic batching
Author: gauravpandeyamu
Pull request: no

We traced the encoder of  LEDForConditionalGeneration using TorchScript and passed a different batch size to the traced encoder as follows.

```
import torch

from transformers import LEDForConditionalGeneration
class WrappedModel(torch.nn.Module):
    def __init__(self):
        super(WrappedModel, self).__init__()
        self.model = LEDForConditionalGeneration.from_pretrained("allenai/led-base-16384", torchscript=True).led.encoder
    def forward(self, data):
        return self.model(data)

example = torch.zeros((1,128), dtype=torch.long)+ 10 # bsz , seqlen
pt_model = WrappedModel().eval()
traced_script_module = torch.jit.trace(pt_model, example)
example_dynamic_batch = torch.zeros((4,128), dtype=torch.long) # bsz , seqlen
traced_script_module(example_dynamic_batch)
```
Being able to vary the batch size during deployment is necessary for dynamic batching to work (for instance, when using Triton inference server). 
Passing a different batch size than the one used during tracing results in the following error:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-4-5e85f8f70ee7> in <module>
     10 pt_model = WrappedModel().eval()
     11 traced_script_module = torch.jit.trace(pt_model, example)
---> 12 traced_script_module(example.repeat(4,1))

/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/transformers/models/led/modeling_led.py(447): _sliding_chunks_query_key_matmul
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/transformers/models/led/modeling_led.py(202): forward
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py(709): _slow_forward
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py(725): _call_impl
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/transformers/models/led/modeling_led.py(725): forward
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py(709): _slow_forward
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py(725): _call_impl
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/transformers/models/led/modeling_led.py(914): forward
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py(709): _slow_forward
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py(725): _call_impl
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/transformers/models/led/modeling_led.py(1838): forward
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py(709): _slow_forward
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py(725): _call_impl
<ipython-input-4-5e85f8f70ee7>(7): forward
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py(709): _slow_forward
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/nn/modules/module.py(725): _call_impl
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/jit/_trace.py(940): trace_module
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/torch/jit/_trace.py(742): trace
<ipython-input-4-5e85f8f70ee7>(11): <module>
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/IPython/core/interactiveshell.py(3441): run_code
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/IPython/core/interactiveshell.py(3361): run_ast_nodes
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/IPython/core/interactiveshell.py(3170): run_cell_async
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/IPython/core/async_helpers.py(68): _pseudo_sync_runner
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/IPython/core/interactiveshell.py(2944): _run_cell
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/IPython/core/interactiveshell.py(2899): run_cell
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/ipykernel/zmqshell.py(539): run_cell
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/ipykernel/ipkernel.py(302): do_execute
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/tornado/gen.py(234): wrapper
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/ipykernel/kernelbase.py(538): execute_request
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/tornado/gen.py(234): wrapper
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/ipykernel/kernelbase.py(261): dispatch_shell
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/tornado/gen.py(234): wrapper
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/ipykernel/kernelbase.py(358): process_one
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/tornado/gen.py(775): run
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/tornado/gen.py(814): inner
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/tornado/ioloop.py(741): _run_callback
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/tornado/ioloop.py(688): <lambda>
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/asyncio/events.py(88): _run
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/asyncio/base_events.py(1758): _run_once
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/asyncio/base_events.py(523): run_forever
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/tornado/platform/asyncio.py(199): start
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/ipykernel/kernelapp.py(619): start
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/traitlets/config/application.py(845): launch_instance
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/site-packages/ipykernel_launcher.py(16): <module>
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/runpy.py(85): _run_code
/dccstor/gpandey11/anaconda3/envs/test/lib/python3.7/runpy.py(193): _run_module_as_main
RuntimeError: shape '[12, 1, 512, 513]' is invalid for input of size 12607488
```

Dynamic batching works fine with the BERT model. For example, the following code gives the correct output.
```
import torch
from transformers import BertForSequenceClassification

class WrappedModel(torch.nn.Module):
    def __init__(self):
        super(WrappedModel, self).__init__()
        self.model = BertForSequenceClassification.from_pretrained('bert-base-uncased', torchscript=True)
    def forward(self, data):
        return self.model(data)

example = torch.zeros((1,128), dtype=torch.long)+ 10 # bsz , seqlen
pt_model = WrappedModel().eval()
traced_script_module = torch.jit.trace(pt_model, example)
example_dynamic_batch = torch.zeros((4,128), dtype=torch.long) # bsz , seqlen
traced_script_module(example_dynamic_batch)
```
## Environment
```
- `transformers` version: 4.3.2
- Platform: Linux-4.18.0-240.22.1.el8_3.x86_64-x86_64-with-redhat-8.3-Ootpa
- Python version: 3.7.0
- PyTorch version (GPU?): 1.7.1 (True)
- Tensorflow version (GPU?): 2.2.0 (False)
- Using GPU in script?: False
- Using distributed or parallel set-up in script?: False
```


Comments:
Hey @gauravpandeyamu, 

This seems to be a rather edge-casy and difficult error to debug! I'm not sure if I manage to have the time to look into it. In a first step could you try to use current master instead of Transformers 4.3.2 to see if this changed anything in the error message?

@patrickvonplaten Sure, I will try it today and let you know. Thanks

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.