diff --git a/cog.yaml b/cog.yaml
index 4e0f631..42e8f9b 100644
--- a/cog.yaml
+++ b/cog.yaml
@@ -7,6 +7,6 @@ build:
     - "libglib2.0-0"
   python_packages:
     - "torch==1.10.1"
-    - "flax==0.4.2"
+    - "flax==0.5.2"
 
 predict: "predict.py:Predictor"