Fix JAX CI Cuda error by fixing JAX version (#19161)
This commit is contained in:
parent
8d58b790e4
commit
830bea69e6
@ -7,7 +7,8 @@ torch>=2.1.0
|
||||
torchvision>=0.16.0
|
||||
|
||||
# Jax with cuda support.
|
||||
# TODO: 0.4.24 has an updated Cuda version breaks Jax CI.
|
||||
--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
|
||||
jax[cuda12_pip]
|
||||
jax[cuda12_pip]==0.4.23
|
||||
|
||||
-r requirements-common.txt
|
||||
|
Loading…
Reference in New Issue
Block a user