From 49b2fc50f0e9a8cd2bbf54c8d4bb5bfc7a546077 Mon Sep 17 00:00:00 2001 From: Akshita Bhagia Date: Mon, 2 Oct 2023 11:11:48 -0700 Subject: [PATCH] higher flax version --- .github/workflows/main.yml | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 82689913..364535c5 100644 --- 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.4.1 jaxlib==0.4.1 tensorflow-cpu==2.9.1 optax==0.1.3 + pip install flax==0.6.1 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' diff --git a/pyproject.toml b/pyproject.toml index 933688f1..e0f401e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -90,7 +90,7 @@ fairscale = [ flax = [ "datasets>=1.12,<3.0", "jax>=0.4.1,<=0.4.13", - "flax>=0.5.0,<=0.7.0", + "flax>=0.6.1,<=0.7.0", "optax>=0.1.2", "tensorflow-cpu>=2.9.1" ]