install torch torchvision and pyro-ppl before jax

QuantEcon · Feb 1, 2024 · 20dd5cf · 20dd5cf
1 parent 4c0fdad
commit 20dd5cf
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/cuda-12.3.1-anaconda-2023-09-py311/Dockerfile b/cuda-12.3.1-anaconda-2023-09-py311/Dockerfile
@@ -42,6 +42,7 @@ ENV PATH /opt/conda/envs/quantecon/bin:$PATH
 
 # Install JAX
 RUN nvcc --version
+RUN pip install torch torchvision pyro-ppl
 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