Tracing LayoutLMv2 results in wrong input_shape dimension
Author: mikkeldenker
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.11.3
- Platform: macOS-12.0-arm64-arm-64bit
- Python version: 3.9.0
- PyTorch version (GPU?): 1.10.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.


Models:

- ALBERT, BERT, XLM, DeBERTa, DeBERTa-v2, ELECTRA, MobileBert, SqueezeBert: @LysandreJik
- encoder-decoder models (For example, BlenderBot, BART, Marian, Pegasus, T5, ByT5): @patrickvonplaten, @patil-suraj
- Longformer, Reformer, TransfoXL, XLNet, FNet: @patrickvonplaten
- FSMT: @stas00
- Funnel: @sgugger
- GPT-2, GPT: @patrickvonplaten, @LysandreJik
- RAG, DPR: @patrickvonplaten, @lhoestq
- TensorFlow: @Rocketknight1
- JAX/Flax: @patil-suraj @patrickvonplaten 
- TAPAS, LayoutLM, LayoutLMv2, LUKE, ViT, BEiT, DEiT, DETR, CANINE: @NielsRogge
- GPT-Neo, GPT-J, CLIP: @patil-suraj
- Wav2Vec2, HuBERT, SpeechEncoderDecoder: @patrickvonplaten, @anton-l

If the model isn't in the list, ping @LysandreJik who will redirect you to the correct contributor.

Library:

- Benchmarks: @patrickvonplaten
- Deepspeed: @stas00
- Ray/raytune: @richardliaw, @amogkam
- Text generation: @patrickvonplaten
- Tokenizers: @LysandreJik
- Trainer: @sgugger
- Pipelines: @Narsil
- Speech: @patrickvonplaten, @anton-l
- Vision: @NielsRogge, @sgugger

Documentation: @sgugger

Model hub:

- for issues with a model, report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj

For research projetcs, please ping the contributor directly. For example, on the following projects:

- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

@NielsRogge 

## Information

When tracing the LayoutLMv2 model, I get the following error:
```
RuntimeError: The expanded size of the tensor (561) must match the existing size (512) at non-singleton dimension 1.  Target sizes: [1, 561].  Tensor sizes: [1, 512]
```

This seems to be caused by [this line](https://github.com/huggingface/transformers/blob/efea0f868bd381244e3cef51b388293e41a36d1e/src/transformers/models/layoutlmv2/modeling_layoutlmv2.py#L862) in the model. I think pytorch might be reusing the memory for `final_shape` and `input_shape` when tracing, so when `final_shape` is updated it also updates `input_shape` which leads to the mismatch in dimensions later on. I've currently solved the problem by changing the shapes to 
```
final_shape = list(torch.empty(size=input_shape).size())
visual_shape = list(torch.empty(size=input_shape).size())
```

I've also changed `visual_shape` for good measure, but it doesn't really seem to be necessary. With the above changes I'm able to successfully trace the model. It seems a bit overkill to allocate an entire new tensor, just to get a copy of the shape but I didn't really find any other solution. Please let me know if there is a better solution.

I'm happy to submit a PR with the fixes if you want.

## To reproduce

Steps to reproduce the behavior:

1. Initialise `LayoutLMv2Model`
2. Trace the model using `torch.jit.trace`

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->


Comments:
Hi,

Thanks for working on making LayoutLMv2 torchscriptable. It would indeed be great if you can open a PR for this.

Also, if LayoutLMv2 works with TorchScript, we can update `tests/test_modeling_layoutlmv2.py`, to also take into account the TorchScript tests (as these are not run right now as seen [here](https://github.com/huggingface/transformers/blob/efea0f868bd381244e3cef51b388293e41a36d1e/tests/test_modeling_layoutlmv2.py#L262)).

That's really the least I can do compared to how you guys make NLP easily accessible to everyone. So thank you for all your amazing work!

I've created a pull request that should fix the issue, and will therefore close this. We can continue the discussion in the PR if needed.