gpt2 error using torch.jit.trace 
## Environment info
- `transformers` version: 4.16.2
- Platform: ubuntu
- Python version: 3.8.12
- PyTorch version (GPU?): 1.10.0a0+0aef44c
- Tensorflow version (GPU?):
- Using GPU in script?:
- Using distributed or parallel set-up in script?:

#### the environment is built from a docker image from `https://github.com/NVIDIA/Torch-TensorRT`, CONTAINER VERSION ==21.10

### Who can help
@patrickvonplaten, @LysandreJik

## To reproduce
issue reported before in #8263 #3954

```
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')
model.eval()

tokens=tokenizer('The cat is on the table.', return_tensors='pt')['input_ids']

with torch.jit.optimized_execution(True):
    traced_model = torch.jit.trace(model, tokens)
```

## information
```
/opt/conda/lib/python3.8/site-packages/transformers/models/gpt2/modeling_gpt2.py:196: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  attn_weights = attn_weights / (float(value.size(-1)) ** 0.5)
Traceback (most recent call last):
  File "generate_speed.py", line 282, in <module>
    tokenizer, eod_id, sep_id, unk_id, model = load_model()
  File "generate_speed.py", line 252, in load_model
    traced_model = torch.jit.trace(model, tokens)
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_trace.py", line 741, in trace
    return trace_module(
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_trace.py", line 958, in trace_module
    module._c._create_method_from_trace(
RuntimeError: Tracer cannot infer type of CausalLMOutputWithCrossAttentions(loss=None, logits=tensor([[[  1.6039, -14.8921, -12.7946,  ..., -12.5804, -13.8293, -13.4610],
         [ -1.9447, -17.5292, -15.7436,  ..., -15.3740, -17.2827, -16.3414],
         [ -2.1778, -15.4163, -13.4095,  ..., -12.6063, -14.2943, -14.2544],
         [ -1.2752, -13.6237, -12.9004,  ..., -12.0724, -11.4089, -11.1407]]],
       device='cuda:0'), past_key_values=((tensor([[[[ 1.3004e+00, -9.0571e-02, -1.1026e+00,  ..., -4.1700e-01,
            2.2270e-01, -1.0456e-02],
          [-1.3599e+00, -9.8881e-01,  3.2034e-01,  ...,  2.1444e-03,
...
[-4.8185e-01,  6.5437e-01, -1.7764e-01,  ...,  1.1841e+00,
            2.4035e-01,  7.0395e-01]]]], grad_fn=<PermuteBackward0>))), hidden_states=None, attentions=None, cross_attentions=None)
:Dictionary inputs to traced functions must have consistent type. Found Tensor and Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor], Tuple[Tensor, Tensor]]
```
dose anyone knows how to solve this? many thanks!<|||||>use `return_dict=False`<|||||>Hello! I recommend reading the following, which mentions what must be done to run models with JIT: https://huggingface.co/docs/transformers/v4.16.2/en/serialization#torchscript<|||||>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.<|||||>try specifying torchscript is True when loading the model
`model = GPT2LMHeadModel.from_pretrained('gpt2',  torchscript=True)`<|||||>I have the same problem, how to solve it?<|||||>Same as me,can anyone solve this problem?<|||||>GPT2LMHeadModel.from_pretrained("gpt2", torchscript=True).eval()