fix jax installation for CUDA12

mlcommons · Feb 29, 2024 · d01c101 · d01c101
1 parent 335acbc
commit d01c101
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/setup.cfg b/setup.cfg
@@ -128,8 +128,7 @@ jax_cpu =
 # JAX GPU
 jax_gpu =
   %(jax_core_deps)s
-  jax[cuda]==0.4.10
-  jaxlib==0.4.10+cuda11.cudnn86
+  jax[cuda12_local]==0.4.10
 
 # PyTorch CPU
 pytorch_cpu =