diff --git a/.github/workflows/publish-example-images.yaml b/.github/workflows/publish-example-images.yaml
index ebf20a721f..8e61ef9c4f 100644
--- a/.github/workflows/publish-example-images.yaml
+++ b/.github/workflows/publish-example-images.yaml
@@ -57,5 +57,7 @@ jobs:
             context: examples/mxnet/tune
           - component-name: pytorch-dist-mnist
             dockerfile: examples/pytorch/mnist/Dockerfile
+            context: examples/pytorch/mnist
           - component-name: pytorch-dist-mnist-mpi
             dockerfile: examples/pytorch/mnist/Dockerfile-mpi
+            context: examples/pytorch/mnist
diff --git a/examples/pytorch/elastic/imagenet/Dockerfile b/examples/pytorch/elastic/imagenet/Dockerfile
index 80c4fe07b9..25ecc49577 100644
--- a/examples/pytorch/elastic/imagenet/Dockerfile
+++ b/examples/pytorch/elastic/imagenet/Dockerfile
@@ -1,3 +1,6 @@
+# We need to use the nvcr.io/nvidia/pytorch image as a base image to support both linux/amd64 and linux_arm64 platforms.
+# 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
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:24.01-py3
 FROM $BASE_IMAGE
 
diff --git a/examples/pytorch/mnist/Dockerfile b/examples/pytorch/mnist/Dockerfile
index ce4e642f72..92e6c5e220 100644
--- a/examples/pytorch/mnist/Dockerfile
+++ b/examples/pytorch/mnist/Dockerfile
@@ -1,3 +1,6 @@
+# We need to use the nvcr.io/nvidia/pytorch image as a base image to support both linux/amd64 and linux_arm64 platforms.
+# 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
 
 RUN pip install tensorboardX==2.6.2
diff --git a/examples/pytorch/mnist/Dockerfile-mpi b/examples/pytorch/mnist/Dockerfile-mpi
index a578373911..57b69ef692 100644
--- a/examples/pytorch/mnist/Dockerfile-mpi
+++ b/examples/pytorch/mnist/Dockerfile-mpi
@@ -1,3 +1,6 @@
+# We need to use the nvcr.io/nvidia/pytorch image as a base image to support both linux/amd64 and linux_arm64 platforms.
+# 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
 
 RUN pip install tensorboardX==2.6.2
diff --git a/examples/pytorch/smoke-dist/Dockerfile b/examples/pytorch/smoke-dist/Dockerfile
index 9b54f4e91c..2760ee2e6f 100644
--- a/examples/pytorch/smoke-dist/Dockerfile
+++ b/examples/pytorch/smoke-dist/Dockerfile
@@ -1,3 +1,6 @@
+# We need to use the nvcr.io/nvidia/pytorch image as a base image to support both linux/amd64 and linux_arm64 platforms.
+# 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
 
 RUN mkdir -p /opt/mlkube