diff --git a/.github/workflows/ut.yml b/.github/workflows/ut.yml
index cb794053d..6fe102d26 100644
--- a/.github/workflows/ut.yml
+++ b/.github/workflows/ut.yml
@@ -22,7 +22,7 @@ jobs:
           conda create -n dmff -y python=${{ matrix.python-version }} numpy openmm==7.7.0 pytest rdkit openbabel mdtraj ambertools -c conda-forge
           conda activate dmff
           pip install --upgrade pip
-          pip install jax jaxlib jaxopt networkx parmed pymbar==4.0.1 optax tqdm
+          pip install jax==0.4.24 jaxlib jaxopt networkx parmed pymbar==4.0.1 optax tqdm
       - name: Install DMFF
         run: |
           source $CONDA/bin/activate dmff && pip install .