diff --git a/tests/notebooks/training/training-integration.ipynb b/tests/notebooks/training/training-integration.ipynb
index 207ede3..1532ee1 100644
--- a/tests/notebooks/training/training-integration.ipynb
+++ b/tests/notebooks/training/training-integration.ipynb
@@ -415,10 +415,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "PYTORCHJOB_NAME = \"pytorch-dist-mnist-gloo\"\n",
+    "PYTORCHJOB_NAME = \"pytorch-mnist-gloo\"\n",
     "PYTORCHJOB_CONTAINER = \"pytorch\"\n",
-    "PYTORCHJOB_IMAGE = \"kubeflow/pytorch-dist-mnist:v1-3a360ba\"\n",
-    "# The image above should be updated with each release with the latest available in the registry."
+    "PYTORCHJOB_IMAGE = \"kubeflowkatib/pytorch-mnist-cpu:v0.16.0\"\n",
+    "# The image above should be updated with each release with the corresponding Katib version used in CKF release.\n",
+    "# Note that instead of using the [image from training-operator repository](https://github.com/kubeflow/training-operator/blob/master/examples/pytorch/mnist/Dockerfile),\n",
+    "# the one [from Katib](https://github.com/kubeflow/katib/blob/master/examples/v1beta1/trial-images/pytorch-mnist/Dockerfile.cpu) is being used\n",
+    "# due to the large size of the first one."
    ]
   },
   {
@@ -430,7 +433,8 @@
     "container = V1Container(\n",
     "    name=PYTORCHJOB_CONTAINER,\n",
     "    image=PYTORCHJOB_IMAGE,\n",
-    "    args=[\"--backend\", \"gloo\"],\n",
+    "    args=[\"--backend\", \"gloo\", \"--epochs\", \"2\"],\n",
+    "    # Passing `epochs`argument since kubeflowkatib image defaults to 10.\n",
     ")\n",
     "\n",
     "replica_spec = V1ReplicaSpec(\n",