diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
index 3d892f8a..55878de4 100644
--- a/Dockerfile.tmpl
+++ b/Dockerfile.tmpl
@@ -38,9 +38,9 @@ RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/lib
 {{ end }}
 
 # Keep these variables in sync if base image is updated.
-ENV TENSORFLOW_VERSION=2.15.0
+ENV TENSORFLOW_VERSION=2.16.1
 # See https://github.com/tensorflow/io#tensorflow-version-compatibility
-ENV TENSORFLOW_IO_VERSION=0.35.0
+ENV TENSORFLOW_IO_VERSION=0.37.0
 
 # We need to redefine the ARG here to get the ARG value defined above the FROM instruction.
 # See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact
@@ -198,11 +198,11 @@ RUN apt-get update && \
 
 RUN pip install -f http://h2o-release.s3.amazonaws.com/h2o/latest_stable_Py.html h2o && /tmp/clean-layer.sh
 
-# b/318672158 Use simply tensorflow-probability once > 0.23.0 is released.
 RUN pip install \
         "tensorflow==${TENSORFLOW_VERSION}" \
         "tensorflow-io==${TENSORFLOW_IO_VERSION}" \
-        git+https://github.com/tensorflow/probability.git@fbc5ebe9b1d343113fb917010096cfd88b32eecf \
+        tensorflow-probability \
+        tensorflow_decision_forests \
         tensorflow_text \
         "tensorflow_hub>=0.16.0" \
         # b/331799280 remove once other packages over to dm-tre
@@ -210,20 +210,13 @@ RUN pip install \
         tf-keras && \
     /tmp/clean-layer.sh
 
-# b/318672158 Use simply tensorflow_decision_forests on next release, expected with tf 2.16
-RUN pip install tensorflow_decision_forests==1.8.1 --no-deps && \
-    /tmp/clean-layer.sh
-
 RUN chmod +x /tmp/keras_patch.sh && \
     /tmp/keras_patch.sh
 
 ADD patches/keras_internal.py /opt/conda/lib/python3.10/site-packages/tensorflow_decision_forests/keras/keras_internal.py
 ADD patches/keras_internal_test.py /opt/conda/lib/python3.10/site-packages/tensorflow_decision_forests/keras/keras_internal_test.py
 
-# Remove "--no-deps" flag and "namex" package once Keras 3.* is included in our base image.
-# We ignore dependencies since tf2.15 and Keras 3.* should work despite pip saying it won't.
-# Currently, keras tries to install a nightly version of tf 2.16: https://github.com/keras-team/keras/blob/fe2f54aa5bc42fb23a96449cf90434ab9bb6a2cd/requirements.txt#L2
-RUN pip install --no-deps "keras>3" keras-cv keras-nlp namex && \
+RUN pip install "keras>3" keras-cv keras-nlp && \
     /tmp/clean-layer.sh
 
 # b/328788268 libpysal 4.10 seems to fail with "module 'shapely' has no attribute 'Geometry'. Did you mean: 'geometry'"
diff --git a/config.txt b/config.txt
index 6afee191..0ff832d6 100644
--- a/config.txt
+++ b/config.txt
@@ -1,7 +1,7 @@
 BASE_IMAGE_REPO=gcr.io/deeplearning-platform-release
-BASE_IMAGE_TAG=m114
-CPU_BASE_IMAGE_NAME=tf2-cpu.2-15.py310
-GPU_BASE_IMAGE_NAME=tf2-gpu.2-15.py310
+BASE_IMAGE_TAG=m122
+CPU_BASE_IMAGE_NAME=tf2-cpu.2-16.py310
+GPU_BASE_IMAGE_NAME=tf2-gpu.2-16.py310
 LIGHTGBM_VERSION=4.2.0
 TORCH_VERSION=2.1.2
 TORCHAUDIO_VERSION=2.1.2
@@ -9,4 +9,4 @@ TORCHTEXT_VERSION=0.16.2
 TORCHVISION_VERSION=0.16.2
 JAX_VERSION=0.4.26
 CUDA_MAJOR_VERSION=12
-CUDA_MINOR_VERSION=1
+CUDA_MINOR_VERSION=3