reverse numpyro and jax installs

QuantEcon · Feb 1, 2024 · 4c0fdad · 4c0fdad
1 parent 4c9af41
commit 4c0fdad
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/cuda-12.3.1-anaconda-2023-09-py311/Dockerfile b/cuda-12.3.1-anaconda-2023-09-py311/Dockerfile
@@ -42,5 +42,6 @@ ENV PATH /opt/conda/envs/quantecon/bin:$PATH
 
 # Install JAX
 RUN nvcc --version
-RUN pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 RUN pip install "numpyro[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+# Note: always install jax[cuda] last to ensure proper CUDA+CUDANN version linking
+RUN pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html