--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pynvml==11.5.0
torch==2.3.0