update orbax dependency due to breaking changes in 0.10.3 (#139)

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