try _pip config for install and show nvidia-smi

QuantEcon · Jan 31, 2024 · 66b5175 · 66b5175
1 parent 0e14d2f
commit 66b5175
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
@@ -41,8 +41,9 @@ RUN echo "source activate quantecon" > ~/.bashrc
 ENV PATH /opt/conda/envs/quantecon/bin:$PATH
 
 # Install JAX
+RUN nvidia-smi
 RUN nvcc --version
 RUN echo "$LD_LIBRARY_PATH"
-RUN pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+RUN pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 COPY check_jax.py .
 RUN python check_jax.py