numpy
git+https://github.com/matthias-wright/flaxmodels.git
tensorflow-datasets
optax
argparse
wandb
tqdm
