Jit.trace not working for ViT model

vision
1

torch.jit.trace is not working for torchvision vision transformer models, i.e. vit_b_16, vit_b_32, and vit_l_16, vit_l_32. It made a TracingCheckError.

2

This is expected for ViT models. Just use jit.trace() instead. For example:
ViT_b_16_script = torch.jit.script(ViT_b_16)