diff --git a/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.cpu b/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.cpu
index 08ddaa1c4a7..bb599e43683 100644
--- a/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.cpu
+++ b/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.cpu
@@ -6,6 +6,7 @@ ADD examples/v1beta1/trial-images/darts-cnn-cifar10 ${TARGET_DIR}
 
 WORKDIR  ${TARGET_DIR}
 
+RUN pip install --prefer-binary --no-cache-dir torch==2.2.1 torchvision==0.17.1
 RUN pip install --prefer-binary --no-cache-dir -r requirements.txt
 RUN chgrp -R 0 ${TARGET_DIR} \
   && chmod -R g+rwX ${TARGET_DIR}
diff --git a/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu b/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu
index 83747235873..ee27092a2ca 100644
--- a/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu
+++ b/examples/v1beta1/trial-images/darts-cnn-cifar10/Dockerfile.gpu
@@ -1,7 +1,7 @@
 # We need to use the nvcr.io/nvidia/pytorch image as a base image to support both linux/amd64 and linux_arm64 platforms.
-# PyTorch=1.13.0, cuda=11.8.0
-# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-22-11.html#rel-22-11
-FROM nvcr.io/nvidia/pytorch:22.11-py3
+# PyTorch=2.2.0, cuda=12.3.2
+# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-01.html#rel-24-01
+FROM nvcr.io/nvidia/pytorch:24.01-py3
 
 ENV TARGET_DIR /opt/darts-cnn-cifar10
 
diff --git a/examples/v1beta1/trial-images/darts-cnn-cifar10/requirements.txt b/examples/v1beta1/trial-images/darts-cnn-cifar10/requirements.txt
index f3f7e5737ff..e20575024df 100644
--- a/examples/v1beta1/trial-images/darts-cnn-cifar10/requirements.txt
+++ b/examples/v1beta1/trial-images/darts-cnn-cifar10/requirements.txt
@@ -1,3 +1 @@
-torch==1.13.1
-torchvision==0.14.1
 Pillow>=9.1.1
diff --git a/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.cpu b/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.cpu
index f1641889fe8..6dca2fa5089 100644
--- a/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.cpu
+++ b/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.cpu
@@ -6,6 +6,7 @@ WORKDIR /opt/pytorch-mnist
 
 # Add folder for the logs.
 RUN mkdir /katib
+RUN pip install --prefer-binary --no-cache-dir torch==2.2.1 torchvision==0.17.1
 RUN pip install --prefer-binary --no-cache-dir -r requirements.txt
 
 RUN chgrp -R 0 /opt/pytorch-mnist \
diff --git a/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu b/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu
index 0216c351606..53f58642e6c 100644
--- a/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu
+++ b/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.gpu
@@ -1,7 +1,7 @@
 # We need to use the nvcr.io/nvidia/pytorch image as a base image to support both linux/amd64 and linux_arm64 platforms.
-# PyTorch=1.13.0, cuda=11.8.0
-# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-22-11.html#rel-22-11
-FROM nvcr.io/nvidia/pytorch:22.11-py3
+# PyTorch=2.2.0, cuda=12.3.2
+# Ref: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-01.html#rel-24-01
+FROM nvcr.io/nvidia/pytorch:24.01-py3
 
 ADD examples/v1beta1/trial-images/pytorch-mnist /opt/pytorch-mnist
 
diff --git a/examples/v1beta1/trial-images/pytorch-mnist/requirements.txt b/examples/v1beta1/trial-images/pytorch-mnist/requirements.txt
index 20af364afd9..23b6c8aaa07 100644
--- a/examples/v1beta1/trial-images/pytorch-mnist/requirements.txt
+++ b/examples/v1beta1/trial-images/pytorch-mnist/requirements.txt
@@ -1,4 +1,2 @@
 cloudml-hypertune==0.1.0.dev6
-torch==1.13.1
-torchvision==0.14.1
 Pillow>=9.1.1
diff --git a/sdk/python/v1beta1/kubeflow/katib/constants/constants.py b/sdk/python/v1beta1/kubeflow/katib/constants/constants.py
index b99de7b98c6..9af281524cd 100644
--- a/sdk/python/v1beta1/kubeflow/katib/constants/constants.py
+++ b/sdk/python/v1beta1/kubeflow/katib/constants/constants.py
@@ -52,7 +52,7 @@
 # TODO (andreyvelich): Implement list_base_images function to get each image description.
 BASE_IMAGE_TENSORFLOW = "docker.io/tensorflow/tensorflow:2.13.0"
 BASE_IMAGE_TENSORFLOW_GPU = "docker.io/tensorflow/tensorflow:2.13.0-gpu"
-BASE_IMAGE_PYTORCH = "docker.io/pytorch/pytorch:1.12.1-cuda11.3-cudnn8-runtime"
+BASE_IMAGE_PYTORCH = "docker.io/pytorch/pytorch:2.2.1-cuda12.1-cudnn8-runtime"
 BASE_IMAGE_MXNET = "docker.io/mxnet/python:1.9.1_native_py3"
 
 DEFAULT_DB_MANAGER_ADDRESS = "katib-db-manager.kubeflow:6789"