Update jax dependencies in GitHub actions as well

allenai · Apr 3, 2024 · f36919b · f36919b
1 parent 1418745
commit f36919b
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.7" "jax<0.5" "jaxlib<0.5" "tensorflow-cpu<2.10" "optax<0.2"
+          pip install flax jax jaxlib "tensorflow-cpu>=2.9.1" optax
 
       - name: Install editable (no cache hit)
         if: steps.virtualenv-cache.outputs.cache-hit != 'true'