torch
torchvision
transformers
jax
git+https://github.com/patil-suraj/vqgan-jax.git