match in github workflow

allenai · Oct 2, 2023 · 0b677cc · 0b677cc
1 parent bea40bf
commit 0b677cc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -211,7 +211,7 @@ jobs:
         if: steps.virtualenv-cache.outputs.cache-hit != 'true' && (contains(matrix.task.extras, 'flax') || contains(matrix.task.extras, 'all'))
         run: |
           . .venv/bin/activate
-          pip install flax==0.5.0 jax==0.3.13 jaxlib==0.3.10 tensorflow-cpu==2.9.1 optax==0.1.3
+          pip install flax==0.5.0 jax==0.4.1 jaxlib==0.4.1 tensorflow-cpu==2.9.1 optax==0.1.3
 
       - name: Install editable (no cache hit)
         if: steps.virtualenv-cache.outputs.cache-hit != 'true'