diff --git a/requirements.txt b/requirements.txt
index 3c07a14..eccd10b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,7 @@
 jax
-orbax
+orbax-export
+orbax-checkpoint
+tensorflow
 seqio
 einops
 absl-py
@@ -9,4 +11,4 @@ tensorboard
 # Following packages are only required for running simply/utils/math_eval.py
 sympy
 pylatexenc
-pytest
\ No newline at end of file
+pytest