export USE_TORCH=False
python3 -m venv jax_env
source ~/jax_env/bin/activate
pip install --upgrade pip
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip install transformers datasets einops flax aiohttp jsonlines zstandard wandb SentencePiece nltk prefetch_generator torch tensorboard tensorflow