-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
jax[cuda]==0.4.4
flax>=0.6.1
optax>=0.1.3
bentoml[grpc,grpc-reflection]
tensorflow
tensorflow-datasets
Pillow
pytest
pytest-asyncio
