jax==0.3.21
jaxlib==0.3.20+cuda11.cudnn82 # If this fails, follow the official instructions: pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
flax==0.5.0
torch==1.10.1
torchtext==0.11.1
torchaudio==0.10.1
torchvision==0.11.2
jupytext==1.13.4
matplotlib==3.5.1
black
flake8
google-cloud-storage
ipython
isort
pre-commit
torchfsdd
tensorflow
tqdm
seaborn
celluloid
wandb
datasets
tqdm
hydra-core
omegaconf