

TITLE
LayoutXLM / LayoutLMv2: error when doing export to TorchScript

COMMENTS
4

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
### System Info

- `transformers` version: 4.30.2
- Platform: Linux-4.15.0-212-generic-x86_64-with-glibc2.27
- Python version: 3.10.11
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- 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?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@NielsRogge 
@ArthurZucker

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Run the following code:
```python
import torch
from transformers import AutoModelForTokenClassification

DEVICE = "cuda"

model = AutoModelForTokenClassification.from_pretrained(
    "microsoft/layoutxlm-base", torchscript=True
).to(DEVICE)

input_ids = torch.tensor([[6, 7, 8, 9, 10]], device=DEVICE)
bboxes = torch.tensor(
    [
        [
            (0.7605, 0.0071, 0.8343, 0.0215),
            (0.834388, 0.0071817, 0.855485, 0.0215431),
            (0.855485, 0.0071817, 0.866033, 0.0215431),
            (0.26377, 0.02427, 0.34743, 0.0421),
            (0.347483, 0.024237, 0.369816, 0.04219),
        ]
    ],
    device=DEVICE,
)
bboxes = (bboxes * 1000).to(int)
image = torch.randn((1, 3, 256, 256)).to(torch.uint8).to(DEVICE)
attention_mask = torch.tensor([[1, 1, 1, 1, 1]], device=DEVICE)

torch.jit.trace(model, [input_ids, bboxes, image, attention_mask]
```

### Expected behavior

I expect to get a traced model as a result of the last line of code. But instead I get huge error traceback saying that "graphs differ across invocations". I see that @NielsRogge already did some changes to the LayoutLMv2 code to fix model tracing. AFAIK LayoutXLM just uses LayoutLMv2 model code under the hood so I was expecting to get a traced model with no problem. 

But it looks like the focus of the change was on different errors (#15254). I haven't found any mentions about my problem anywhere here except for the #17476 where it helped to just disable trace checking. If I try to disable trace checking, the model works at the first inference but predictions start to deviate seriously after that. Every prediction on the same image is *very* different. So I guess it's not the solution in my case. 

Am I doing something wrong here or is this model not really compatible with PyTorch tracing functionality? I'm pretty carefully following the official guide about exporting the model to TorchScript.