use pip for jax

QuantEcon · May 27, 2024 · 95dd642 · 95dd642
1 parent 4ad93be
commit 95dd642
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -22,12 +22,12 @@ jobs:
       - name: Install CUDA
         shell: bash -l {0}
         run: |
-          conda install anaconda::cudatoolkit -y
+          conda install cudatoolkit cuda-nvcc -c nvidia -c anaconda -y
       - name: Install JAX[CUDA] and Numpyro[CUDA]
         shell: bash -l {0}
         run: |
-          conda install jaxlib=*=*cuda* jax cuda-nvcc -c conda-forge -c nvidia -y
-          conda install conda-forge::numpyro -y
+          pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+          pip install --upgrade "numpyro[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
       - name: Check nvidia drivers
         shell: bash -l {0}
         run: |