diff --git a/requirements.txt b/requirements.txt
index bdbebc9..4407704 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 flake8>=6.1.0
-flax==1.3.5
+flax==0.8.3
 jax==0.4.28
 jaxlib==0.4.28
 matplotlib>=3.7.2