diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index fb0647b..fb5e8c7 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -43,7 +43,7 @@ jobs:
       with:
         pre-build-command: |
           apt-get --allow-releaseinfo-change update -y && apt-get install -y git && git --version
-          python -m pip install --upgrade pip setuptools wheel
+          python -m pip install --upgrade pip setuptools wheel "jax[cpu]"
           python -m pip install -e .[composable,neural,docs]
         docs-folder: "docs/"
     - name: Publish documentation