diff --git a/notebooks/responsible_ai/explainable_ai/labs/xai_image_vertex.ipynb b/notebooks/responsible_ai/explainable_ai/labs/xai_image_vertex.ipynb
index 94e67b5b..e3206cd3 100644
--- a/notebooks/responsible_ai/explainable_ai/labs/xai_image_vertex.ipynb
+++ b/notebooks/responsible_ai/explainable_ai/labs/xai_image_vertex.ipynb
@@ -646,8 +646,7 @@
     "    decoded = tf.io.decode_jpeg(bytes_input, channels=3)\n",
     "    decoded = tf.image.convert_image_dtype(decoded, tf.float32)\n",
     "    resized = #TODO: Resize decoded image\n",
-    "    rescale = #TODO: Rescale image\n",
-    "    return rescale\n",
+    "    return resized\n",
     "\n",
     "\n",
     "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
diff --git a/notebooks/responsible_ai/explainable_ai/solutions/xai_image_vertex.ipynb b/notebooks/responsible_ai/explainable_ai/solutions/xai_image_vertex.ipynb
index 2eb4f0e0..fa624f52 100644
--- a/notebooks/responsible_ai/explainable_ai/solutions/xai_image_vertex.ipynb
+++ b/notebooks/responsible_ai/explainable_ai/solutions/xai_image_vertex.ipynb
@@ -646,8 +646,7 @@
     "    decoded = tf.io.decode_jpeg(bytes_input, channels=3)\n",
     "    decoded = tf.image.convert_image_dtype(decoded, tf.float32)\n",
     "    resized = tf.image.resize(decoded, size=(192, 192))\n",
-    "    rescale = tf.cast(resized / 255.0, tf.float32)\n",
-    "    return rescale\n",
+    "    return resized\n",
     "\n",
     "\n",
     "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",