Implementation of Vision Transformer as proposed in the paper https://arxiv.org/abs/2010.11929
This implementation is based on the official jax/flax based implementation. https://github.com/google-research/vision_transformer
The Given Notebook contians the complete code for training ViT model. It includes a data pipeline and code for training on TPU and GPU