diff --git a/requirements.txt b/requirements.txt
index 82ade3e..2c7b4bd 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -24,6 +24,6 @@ tensorflow-datasets>=4.9.6
 ruff>=0.1.5,<=0.2
 git+https://github.com/mlperf/logging.git
 opencv-python-headless==4.10.0.84
-orbax-checkpoint>=0.5.20
+orbax-checkpoint==0.10.2
 tokenizers==0.20.0
 huggingface_hub==0.24.7
\ No newline at end of file