diff --git a/.github/container/Dockerfile.maxtext b/.github/container/Dockerfile.maxtext
index a6338c11d..b5013f3c9 100644
--- a/.github/container/Dockerfile.maxtext
+++ b/.github/container/Dockerfile.maxtext
@@ -67,6 +67,7 @@ for pattern in \
     "s|absl-py|absl-py>=2.1.0|g" \
     "s|protobuf==3.20.3|protobuf>=3.19.0|g" \
     "s|tensorflow-datasets|tensorflow-datasets>=4.8.0|g" \
+    "s|grain-nightly|grain|g" \
   ; do
     sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/requirements.txt;
 done
diff --git a/.github/container/Dockerfile.t5x b/.github/container/Dockerfile.t5x
index ea4bbf2ec..84004ba09 100644
--- a/.github/container/Dockerfile.t5x
+++ b/.github/container/Dockerfile.t5x
@@ -73,6 +73,7 @@ echo "seqio-nightly>=0.0.18.dev20240714" >> /opt/pip-tools.d/requirements-t5x.in
 # 2. Remove head-of-tree specs from select dependencies
 pushd ${SRC_PATH_T5X}
 sed -i "s| @ git+https://github.com/google/flax#egg=flax||g" setup.py
+sed -i "s| @ git+https://github.com/deepmind/optax#egg=optax||g" setup.py
 
 # for ARM64 build
 if [[ "$(dpkg --print-architecture)" == "arm64" ]]; then