Better Jax installation instructions for benchmarking.

axch · Oct 17, 2022 · 5145c5d · 5145c5d
1 parent d116661
commit 5145c5d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/.github/workflows/bench.yaml b/.github/workflows/bench.yaml
@@ -29,7 +29,7 @@ jobs:
     - name: Install system dependencies
       run: |
         ${{ matrix.install_deps }}
-        pip install numpy jax
+        pip install numpy "jax[cpu]"
         echo "${{ matrix.path_extension }}" >> $GITHUB_PATH
 
     - name: Cache