From 1feb3176ee669333f310dd88cc0af5f828600736 Mon Sep 17 00:00:00 2001
From: michalk8 <46717574+michalk8@users.noreply.github.com>
Date: Wed, 29 Mar 2023 16:35:25 +0200
Subject: [PATCH] Pin `chex` version in GPU CI (#346)

* Pin `chex` version

* Fix typo
---
 .github/workflows/tests.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 405b32978..b139d48df 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -55,11 +55,12 @@ jobs:
 
     - name: Install dependencies
       # `jax[cuda]<0.4` because of Docker issues: https://github.com/google/jax/issues/13758
+      # `chex<0.1.7` because it requires `jax>=0.4.6`
       # `flax<0.6.5` because it requires `jax>=0.4.2`
       run: |
         python3 -m pip install --upgrade pip
         python3 -m pip install -e".[test]"
-        python3 -m pip install "flax<0.6.5"
+        python3 -m pip install "flax<0.6.5" "chex<0.1.7"
         python3 -m pip install "jax[cuda]<0.4" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
     - name: Nvidia SMI