Code is inspired by the amazing collection in vit-pytorch and the original Jax repository
The data folder only contains a tiny example of 2k images for training and validation, please refer to the TopTagging dataset TopTagging dataset for training on the full 1.2M data set.
We will hopefully provide a notebook to illustrate the technique