diff --git a/.github/workflows/bench.yaml b/.github/workflows/bench.yaml
index f986d4099..83e4f870e 100644
--- 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