diff --git a/requirements.txt b/requirements.txt
index 2505a54d69..f2971a4564 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,9 +1,11 @@
-orbax-checkpoint
+jax>=0.4.23
+jaxlib>=0.4.23
+orbax-checkpoint>=0.5.2
 absl-py
 aqtp
 cloud-tpu-diagnostics
 google-cloud-storage
-flax
+flax>=0.8.0
 ml-collections
 numpy
 optax
@@ -12,8 +14,8 @@ pylint
 pytest
 pytype
 sentencepiece==0.1.97
+tensorflow-text>=2.13.0
+tensorflow>=2.13.0
 tensorflow-datasets
-tensorflow-text
 tensorboardx
-tensorflow
-tensorboard-plugin-profile
+tensorboard-plugin-profile
\ No newline at end of file
diff --git a/setup.sh b/setup.sh
index 2b63de9b63..608f86d1d0 100644
--- a/setup.sh
+++ b/setup.sh
@@ -194,4 +194,4 @@ else
 fi
 
 # Install dependencies from requirements.txt
-cd $run_name_folder_path && pip3 install -r requirements.txt
+cd $run_name_folder_path && pip install --upgrade pip &&  pip3 install -r requirements.txt