torch.fx not working on ViT model
Author: dwlim-nota
Pull request: no
Labels: name: bug

### System Info

transformer version: 4.23.0.dev0
platform: windows 11, AMD64
python version: 3.7.9

### Who can help?

@NielsRogge

### Information

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

### Tasks

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

### Reproduction

1. write below code and execute
```
import torch
import numpy as np
from transformers import ViTFeatureExtractor, ViTModel
from datasets import load_dataset
from torch.fx import symbolic_trace


dataset = load_dataset("huggingface/cats-image")
image = dataset["test"]["image"][0]

feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
model = ViTModel.from_pretrained("google/vit-base-patch16-224-in21k")


img = torch.Tensor(np.random.randn(1, 3, 224, 224))

tracing_info = {"head_mask": None, "interpolate_pos_encoding": None, "bool_masked_pos": None, "output_hidden_states": None, "output_attentions": None, "return_dict": None}
traced = symbolic_trace(model, tracing_info) # bug here

with torch.no_grad():
    outputs = model(img)
    outputs2 = traced(img)
assert torch.allclose(dict(outputs)["last_hidden_state"], outputs2["last_hidden_state"])
```


2. error traceback
```
symbolically traced variables cannot be used as inputs to control flow
  File "C:\Users\NOTA2001\Desktop\abab\transformers\src\transformers\models\vit\modeling_vit.py", line 166, in forward
    if num_channels != self.num_channels:
  File "C:\Users\NOTA2001\Desktop\abab\transformers\src\transformers\models\vit\modeling_vit.py", line 118, in forward
    embeddings = self.patch_embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
  File "C:\Users\NOTA2001\Desktop\abab\transformers\src\transformers\models\vit\modeling_vit.py", line 558, in forward
    pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding
  File "C:\Users\NOTA2001\Desktop\abab\transformers\do_something.py", line 19, in <module>
    traced = symbolic_trace(model, inputs2)
```

### Expected behavior

when this issue was fixed then below code will work

```
import torch
import numpy as np
from transformers import ViTFeatureExtractor, ViTModel
from datasets import load_dataset
from torch.fx import symbolic_trace


dataset = load_dataset("huggingface/cats-image")
image = dataset["test"]["image"][0]

feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
model = ViTModel.from_pretrained("google/vit-base-patch16-224-in21k")


img = torch.Tensor(np.random.randn(1, 3, 224, 224))

tracing_info = {"head_mask": None, "interpolate_pos_encoding": None, "bool_masked_pos": None, "output_hidden_states": None, "output_attentions": None, "return_dict": None}
traced = symbolic_trace(model, tracing_info) # bug here

with torch.no_grad():
    outputs = model(img)
    outputs2 = traced(img)
assert torch.allclose(dict(outputs)["last_hidden_state"], outputs2["last_hidden_state"])
```

Comments:
I am a bit surprised with this situation, as `ViTModel` is tested against `test_torch_fx_xxx` methods, which go through
https://github.com/huggingface/transformers/blob/2d956958252617a178a68a06582c99b133fe7d3d/tests/test_modeling_common.py#L772

However, I realized that `symbolic_trace` is `from transformers.utils.fx import symbolic_trace`
https://github.com/huggingface/transformers/blob/2d956958252617a178a68a06582c99b133fe7d3d/src/transformers/utils/fx.py#L1107
and not `from torch.fx import symbolic_trace`.

Just a remark: It looks like, (if we want) to fully support `from torch.fx import symbolic_trace` in the library, there are more places to be changed.

Thank you for your comment @ydshieh 
I didn't know the fact that `from transformers.utils.fx import symbolic_trace` exists.
Which do you think better, `from torch.fx import symbolic_trace` or `from transformers.utils.fx import symbolic_trace`?
because I am a newbie contributor to transformers, so I want to hear.

If you think it is better to use `from transformers.utils.fx import symbolic_trace`, then I want to close this issue.

Hi @dwlim-nota, 
We actually do support `torch.fx` symbolic tracing though our custom tracer, which is supposed to handle the ViT case and what you do in your PR.

```python
from transformers.utils.fx import symbolic_trace
traced = symbolic_trace(vit_model, input_names=[ "pixel_values"])
```

Compared to the original `torch.fx.symbolic_trace` function, you need to specify which inputs the traced model must have (because our models usually support different inputs, which is not possible in fx), so that is why I specified `pixel_values` here.

thank you for your comment @michaelbenayoun 
I will close this issue :)