Repository: pytorch
Issue: Cannot run FX tracer on a vision transformer model
Labels: triaged,Stale,module: fx

### 🐛 Describe the bug

FX is unable to trace the graph of a vision transformer model:

```python
import torch
import torchvision

# load a vision transformer model from torch.hub
dino = torch.hub.load('facebookresearch/dino:main', 'dino_vits16')
# try and trace the model
train_nodes, eval_nodes = torchvision.models.feature_extraction.get_graph_node_names(dino)
```

This raises the following error:
```
----> 1 train_nodes, eval_nodes = torchvision.models.feature_extraction.get_graph_node_names(dino)

~/.conda/envs/shai-py37/lib/python3.9/site-packages/torchvision/models/feature_extraction.py in get_graph_node_names(model, tracer_kwargs, suppress_diff_warning)
    212     is_training = model.training
    213     train_tracer = NodePathTracer(**tracer_kwargs)
--> 214     train_tracer.trace(model.train())
    215     eval_tracer = NodePathTracer(**tracer_kwargs)
    216     eval_tracer.trace(model.eval())

~/.conda/envs/shai-py37/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py in trace(self, root, concrete_args)
    613                 for module in self._autowrap_search:
    614                     _autowrap_check(patcher, module.__dict__, self._autowrap_function_ids)
--> 615                 self.create_node('output', 'output', (self.create_arg(fn(*args)),), {},
    616                                  type_expr=fn.__annotations__.get('return', None))
    617

~/.cache/torch/hub/facebookresearch_dino_main/vision_transformer.py in forward(self, x)
    208
    209     def forward(self, x):
--> 210         x = self.prepare_tokens(x)
    211         for blk in self.blocks:
    212             x = blk(x)

~/.cache/torch/hub/facebookresearch_dino_main/vision_transformer.py in prepare_tokens(self, x)
    203
    204         # add positional encoding to each token
--> 205         x = x + self.interpolate_pos_encoding(x, w, h)
    206
    207         return self.pos_drop(x)

~/.cache/torch/hub/facebookresearch_dino_main/vision_transformer.py in interpolate_pos_encoding(self, x, w, h)
    175         npatch = x.shape[1] - 1
    176         N = self.pos_embed.shape[1] - 1
--> 177         if npatch == N and w == h:
    178             return self.pos_embed
    179         class_pos_embed = self.pos_embed[:, 0]

~/.conda/envs/shai-py37/lib/python3.9/site-packages/torch/fx/proxy.py in __bool__(self)
    249
    250     def __bool__(self) -> bool:
--> 251         return self.tracer.to_bool(self)
    252
    253     @compatibility(is_backward_compatible=True)

~/.conda/envs/shai-py37/lib/python3.9/site-packages/torch/fx/proxy.py in to_bool(self, obj)
    150         information to the graph node using create_node and can choose to return a value.
    151         """
--> 152         raise TraceError('symbolically traced variables cannot be used as inputs to control flow')
    153
    154     @compatibility(is_backward_compatible=True)

TraceError: symbolically traced variables cannot be used as inputs to control flow
```

This issue might be related to #67970



### Versions

```
PyTorch version: 1.10.1
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 7.8.2003 (Core) (x86_64)
GCC version: (GCC) 8.3.0
Clang version: Could not collect
CMake version: version 2.8.12.2
Libc version: glibc-2.17

Python version: 3.9.7 (default, Sep 16 2021, 13:09:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-3.10.0-1127.19.1.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: False
CUDA runtime version: 10.2.89
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.2
[pip3] torch==1.10.1
[pip3] torchaudio==0.10.1
[pip3] torchtext==0.11.1
[pip3] torchvision==0.11.2
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               10.2.89              hfd86e86_1
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-service               2.4.0            py39h7f8727e_0
[conda] mkl_fft                   1.3.1            py39hd3c417c_0
[conda] mkl_random                1.2.2            py39h51133e4_0
[conda] numpy                     1.21.2           py39h20f2e39_0
[conda] numpy-base                1.21.2           py39h79a1101_0
[conda] pytorch                   1.10.1          py3.9_cuda10.2_cudnn7.6.5_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.10.1               py39_cu102    pytorch
[conda] torchtext                 0.11.1                     py39    pytorch
[conda] torchvision               0.11.2               py39_cu102    pytorch
```

cc @ezyang