This repository includes
Create a virtual env and install requirements.txt.
requirements.txt
pip install -r requirements.txt
For Google Cloud TPU VM please follow necessary installation steps here:
Pytorch on TPU VM JAX/Flax on TPU VM