use conda

.github/workflows/ci.yml

@@ -3,9 +3,9 @@ on: [pull_request]
 jobs:
   preview:
     runs-on: ubuntu-latest-gpu
-    container:
-      image: docker://nvidia/cuda:12.4.1-devel-ubuntu22.04
-      options: --gpus all
+    # container:
+    #   image: docker://nvidia/cuda:12.4.1-devel-ubuntu22.04
+    #   options: --gpus all
     steps:
       - uses: actions/checkout@v4
         with:
@@ -18,12 +18,16 @@ jobs:
           miniconda-version: 'latest'
           python-version: "3.11"
           environment-file: environment.yml
-          activate-environment: quantecon
+          activate-environment: quantecon       
+      - name: Install CUDA
+        shell: bash -l {0}
+        run: |
+          conda install anaconda::cudatoolkit -y
       - name: Install JAX[CUDA] and Numpyro[CUDA]
         shell: bash -l {0}
         run: |
-          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
+          conda install jaxlib=*=*cuda* jax cuda-nvcc -c conda-forge -c nvidia -y
+          conda install conda-forge::numpyro -y
       - name: Check nvidia drivers
         shell: bash -l {0}
         run: |