The provided colab link does not run.

[swyoon]

The authors said the provided google colab link should reproduce the result. However, the colab notebook, when run in a default colab environment, fails even before loading a model.

The main error occurred in this cell:

# Import files from repository.

import sys
if './vision_transformer' not in sys.path:
  sys.path.append('./vision_transformer')

%load_ext autoreload
%autoreload 2

from vit_jax import checkpoint
from vit_jax import models
from vit_jax import train
from vit_jax.configs import augreg as augreg_config
from vit_jax.configs import models as models_config

The error message says

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
[<ipython-input-8-e79c46bdde58>](https://localhost:8080/#) in <module>()
      8 get_ipython().magic('autoreload 2')
      9 
---> 10 from vit_jax import checkpoint
     11 from vit_jax import models
     12 from vit_jax import train

[/content/vision_transformer/vit_jax/checkpoint.py](https://localhost:8080/#) in <module>()
     17 
     18 from absl import logging
---> 19 import flax
     20 from  flax.training import checkpoints
     21 import jax.numpy as jnp

ModuleNotFoundError: No module named 'flax'

---------------------------------------------------------------------------
NOTE: If your import is failing due to a missing package, you can
manually install dependencies using either !pip or !apt.

To view examples of installing some common dependencies, click the
"Open Examples" button below.
---------------------------------------------------------------------------

I wish the authors could look into the issue and keep the work reproducible.