From 8483efa843522dd2b8adefa5af66b49a65251e4c Mon Sep 17 00:00:00 2001
From: Masahiro Hiramori <contact@mshr-h.com>
Date: Wed, 2 Oct 2024 00:18:51 +0900
Subject: [PATCH] use flax 0.8.5

---
 docker/install/ubuntu_install_jax.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/install/ubuntu_install_jax.sh b/docker/install/ubuntu_install_jax.sh
index 2188e5490d48..17114e0efce8 100644
--- a/docker/install/ubuntu_install_jax.sh
+++ b/docker/install/ubuntu_install_jax.sh
@@ -34,4 +34,4 @@ else
 fi
 
 # Install flax
-pip3 install flax~=0.9.0
+pip3 install flax~=0.8.5