torch.jit._trace.TracingCheckError: Tracing failed sanity checks! ERROR: Graphs differed across invocations! 
### System Info

- `transformers` version: 4.31.0
- Platform: Linux-4.18.0-2.4.3.3.kwai.x86_64-x86_64-with-glibc2.17
- Python version: 3.11.4
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- Accelerate version: not installed
- Accelerate config: not found
- 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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

_No response_

### 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

```python
import os
import torch  
from transformers import SwinForImageClassification, TrainingArguments, Trainer

label2id ={'bad': 1, 'good': 0}
id2label = {1:'bad', 0:'good'}

model_name = 'microsoft/swin-base-patch4-window12-384-in22k'
model = SwinForImageClassification.from_pretrained(
    model_name,
    label2id=label2id,
    id2label=id2label,
    ignore_mismatched_sizes=True, 
    torchscript=True
)
'''
ckpt = "./pytorch_model_swin.bin" 
checkpoint = torch.load(ckpt)
model.load_state_dict(checkpoint) # my weight can be successfully loaded into the model
'''
t = torch.randn(1,3,384,384)
t_model = torch.jit.trace(model, t) # error occurred here
```
**Log Message:**
```
---------------------------------------------------------------------------
TracingCheckError                         Traceback (most recent call last)
Cell In[24], line 4
      2 out = model(t)
      3 print(out[0].shape)
----> 4 t_model = torch.jit.trace(model, [t])
      5 print("ok")

File /home/web_server/anaconda3/envs/ram/lib/python3.11/site-packages/torch/jit/_trace.py:794, in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit, example_kwarg_inputs, _store_inputs)
    792         else:
    793             raise RuntimeError("example_kwarg_inputs should be a dict")
--> 794     return trace_module(
    795         func,
    796         {"forward": example_inputs},
    797         None,
    798         check_trace,
    799         wrap_check_inputs(check_inputs),
    800         check_tolerance,
    801         strict,
    802         _force_outplace,
    803         _module_class,
    804         example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),
    805         _store_inputs=_store_inputs
    806     )
    807 if (
    808     hasattr(func, "__self__")
    809     and isinstance(func.__self__, torch.nn.Module)
    810     and func.__name__ == "forward"
    811 ):
    812     if example_inputs is None:

File /home/web_server/anaconda3/envs/ram/lib/python3.11/site-packages/torch/jit/_trace.py:1084, in trace_module(mod, inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit, example_inputs_is_kwarg, _store_inputs)
   1072                 _check_trace(
   1073                     check_inputs,
   1074                     func,
   (...)
   1081                     example_inputs_is_kwarg=example_inputs_is_kwarg,
   1082                 )
   1083             else:
-> 1084                 _check_trace(
   1085                     [inputs],
   1086                     func,
   1087                     check_trace_method,
   1088                     check_tolerance,
   1089                     strict,
   1090                     _force_outplace,
   1091                     True,
   1092                     _module_class,
   1093                     example_inputs_is_kwarg=example_inputs_is_kwarg,
   1094                 )
   1095 finally:
   1096     torch.jit._trace._trace_module_map = old_module_map

File /home/web_server/anaconda3/envs/ram/lib/python3.11/site-packages/torch/utils/_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    112 @functools.wraps(func)
    113 def decorate_context(*args, **kwargs):
    114     with ctx_factory():
--> 115         return func(*args, **kwargs)

File /home/web_server/anaconda3/envs/ram/lib/python3.11/site-packages/torch/jit/_trace.py:562, in _check_trace(check_inputs, func, traced_func, check_tolerance, strict, force_outplace, is_trace_module, _module_class, example_inputs_is_kwarg)
    560 diag_info = graph_diagnostic_info()
    561 if any(info is not None for info in diag_info):
--> 562     raise TracingCheckError(*diag_info)

TracingCheckError: Tracing failed sanity checks!
ERROR: Graphs differed across invocations!
	Graph diff:
		  graph(%self.1 : __torch__.transformers.models.swin.modeling_swin.SwinForImageClassification,
		        %pixel_values : Tensor):
		    %classifier : __torch__.torch.nn.modules.linear.Linear = prim::GetAttr[name="classifier"](%self.1)
		    %swin : __torch__.transformers.models.swin.modeling_swin.SwinModel = prim::GetAttr[name="swin"](%self.1)
...
```

### Expected behavior

Hello, an error occured in the following code when I was using torch.jit.trace to transfer the Swin Transformer model to TorchScript. What should I do to fix it?
Hi it seems you code snippet is somehow different than what has been shown in the log where it has

```Cell In[24], line 4
      2 out = model(t)
      3 print(out[0].shape)
----> 4 t_model = torch.jit.trace(model, [t])
      5 print("ok")
```
Also the code snippet won't work for us as we don't have

ckpt = "./pytorch_model_swin.bin" and `t_model` is defined.

Could you update the code snippet so we can reproduce the error directly? Thanks in advance!

<|||||>Thanks for reply! The code and log message are update as follows. I didn't paste the log message after the 'graph diff' because it's very long.


**Code**
```python
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
import torch  
from transformers import SwinForImageClassification, TrainingArguments, Trainer
label2id ={'bad': 1, 'good': 0}
id2label = {1:'bad', 0:'good'}

model_name = 'microsoft/swin-base-patch4-window12-384-in22k'
model = SwinForImageClassification.from_pretrained(
    model_name,
    label2id=label2id,
    id2label=id2label,
    ignore_mismatched_sizes=True,  # provide this in case you're planning to fine-tune an already fine-tuned checkpoint
    torchscript=True
)

dummy_input = torch.randn(1,3,384,384)
traced_model = torch.jit.trace(model, dummy_input)
print("ok")
```
**Log Message**
```python
'(MaxRetryError("HTTPSConnectionPool(host='huggingface.co', port=443): Max retries exceeded with url: /microsoft/swin-base-patch4-window12-384-in22k/resolve/main/config.json (Caused by ConnectTimeoutError(<urllib3.connection.HTTPSConnection object at 0x7f117efe3690>, 'Connection to huggingface.co timed out. (connect timeout=10)'))"), '(Request ID: 46579f44-3f58-4627-8466-22306063fb52)')' thrown while requesting HEAD https://huggingface.co/microsoft/swin-base-patch4-window12-384-in22k/resolve/main/config.json
Some weights of SwinForImageClassification were not initialized from the model checkpoint at microsoft/swin-base-patch4-window12-384-in22k and are newly initialized because the shapes did not match:
- classifier.weight: found shape torch.Size([21841, 1024]) in the checkpoint and torch.Size([2, 1024]) in the model instantiated
- classifier.bias: found shape torch.Size([21841]) in the checkpoint and torch.Size([2]) in the model instantiated
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
---------------------------------------------------------------------------
TracingCheckError                         Traceback (most recent call last)
Cell In[2], line 18
      9 model = SwinForImageClassification.from_pretrained(
     10     model_name,
     11     label2id=label2id,
   (...)
     14     torchscript=True
     15 )
     17 dummy_input = torch.randn(1,3,384,384)
---> 18 traced_model = torch.jit.trace(model, dummy_input)
     19 print("ok")

File /home/web_server/anaconda3/envs/ram/lib/python3.11/site-packages/torch/jit/_trace.py:794, in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit, example_kwarg_inputs, _store_inputs)
    792         else:
    793             raise RuntimeError("example_kwarg_inputs should be a dict")
--> 794     return trace_module(
    795         func,
    796         {"forward": example_inputs},
    797         None,
    798         check_trace,
    799         wrap_check_inputs(check_inputs),
    800         check_tolerance,
    801         strict,
    802         _force_outplace,
    803         _module_class,
    804         example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),
    805         _store_inputs=_store_inputs
    806     )
    807 if (
    808     hasattr(func, "__self__")
    809     and isinstance(func.__self__, torch.nn.Module)
    810     and func.__name__ == "forward"
    811 ):
    812     if example_inputs is None:

File /home/web_server/anaconda3/envs/ram/lib/python3.11/site-packages/torch/jit/_trace.py:1084, in trace_module(mod, inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit, example_inputs_is_kwarg, _store_inputs)
   1072                 _check_trace(
   1073                     check_inputs,
   1074                     func,
   (...)
   1081                     example_inputs_is_kwarg=example_inputs_is_kwarg,
   1082                 )
   1083             else:
-> 1084                 _check_trace(
   1085                     [inputs],
   1086                     func,
   1087                     check_trace_method,
   1088                     check_tolerance,
   1089                     strict,
   1090                     _force_outplace,
   1091                     True,
   1092                     _module_class,
   1093                     example_inputs_is_kwarg=example_inputs_is_kwarg,
   1094                 )
   1095 finally:
   1096     torch.jit._trace._trace_module_map = old_module_map

File /home/web_server/anaconda3/envs/ram/lib/python3.11/site-packages/torch/utils/_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    112 @functools.wraps(func)
    113 def decorate_context(*args, **kwargs):
    114     with ctx_factory():
--> 115         return func(*args, **kwargs)

File /home/web_server/anaconda3/envs/ram/lib/python3.11/site-packages/torch/jit/_trace.py:562, in _check_trace(check_inputs, func, traced_func, check_tolerance, strict, force_outplace, is_trace_module, _module_class, example_inputs_is_kwarg)
    560 diag_info = graph_diagnostic_info()
    561 if any(info is not None for info in diag_info):
--> 562     raise TracingCheckError(*diag_info)

TracingCheckError: Tracing failed sanity checks!
ERROR: Graphs differed across invocations!
	Graph diff:
...
```<|||||>Thank you for updating, very nice 🤗 

<|||||>Confirmed the reproduction (and yes the log after Graph diff is super super long 😅 <|||||>@fxmarty Could you help on this? You can check the following code snippet.

Basically, it will pass or fail depending on different config values. There is some data flow (tensor bool values) issue in the modeling code, but I am really bad on identifying where is the root cause and how to fix things here.

If possible, that would be great if you can share how you debug this kind of thing with us 🙏 Thank you in advance.
 

```python
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
import torch
from transformers import SwinForImageClassification, TrainingArguments, Trainer, SwinConfig

label2id ={'bad': 1, 'good': 0}
id2label = {1:'bad', 0:'good'}


# this fails
USE_SMALL_CONFIG = False
# this works
# USE_SMALL_CONFIG = True

model_name = 'microsoft/swin-base-patch4-window12-384-in22k'
config = SwinConfig.from_pretrained(model_name)
config.torchscript = True
config.label2id=label2id
config.id2label=id2label

if USE_SMALL_CONFIG:
    config.image_size = 32
    config.patch_size = 2
    config.depths=[1, 2, 1]
    config.num_heads=[2, 2, 4]
    config.window_size=2

model = SwinForImageClassification(config)
dummy_input = torch.randn(1,3, config.image_size, config.image_size)

# make sure it can run in normal mode
o = model(dummy_input)
print("model forward ok")
# trace it
traced_model = torch.jit.trace(model, dummy_input)
print("trace ok")
```