diff --git a/guides/ipynb/keras_hub/semantic_segmentation_deeplab_v3.ipynb b/guides/ipynb/keras_hub/semantic_segmentation_deeplab_v3.ipynb
index 06d4b40a9e..7b366f7fe4 100644
--- a/guides/ipynb/keras_hub/semantic_segmentation_deeplab_v3.ipynb
+++ b/guides/ipynb/keras_hub/semantic_segmentation_deeplab_v3.ipynb
@@ -1053,7 +1053,7 @@
     "test_ds = load(split=\"sbd_eval\")\n",
     "test_ds = preprocess_inputs(test_ds)\n",
     "\n",
-    "images, masks = next(iter(train_ds.take(1)))\n",
+    "images, masks = next(iter(test_ds.take(1)))\n",
     "images = ops.convert_to_tensor(images)\n",
     "masks = ops.convert_to_tensor(masks)\n",
     "preds = ops.expand_dims(ops.argmax(model.predict(images), axis=-1), axis=-1)\n",
diff --git a/guides/keras_hub/semantic_segmentation_deeplab_v3.py b/guides/keras_hub/semantic_segmentation_deeplab_v3.py
index 2ddae379a8..b72541f9a1 100644
--- a/guides/keras_hub/semantic_segmentation_deeplab_v3.py
+++ b/guides/keras_hub/semantic_segmentation_deeplab_v3.py
@@ -827,7 +827,7 @@ def dict_to_tuple(x):
 test_ds = load(split="sbd_eval")
 test_ds = preprocess_inputs(test_ds)
 
-images, masks = next(iter(train_ds.take(1)))
+images, masks = next(iter(test_ds.take(1)))
 images = ops.convert_to_tensor(images)
 masks = ops.convert_to_tensor(masks)
 preds = ops.expand_dims(ops.argmax(model.predict(images), axis=-1), axis=-1)
diff --git a/guides/md/keras_hub/semantic_segmentation_deeplab_v3.md b/guides/md/keras_hub/semantic_segmentation_deeplab_v3.md
index d08679b536..78fbe6b89c 100644
--- a/guides/md/keras_hub/semantic_segmentation_deeplab_v3.md
+++ b/guides/md/keras_hub/semantic_segmentation_deeplab_v3.md
@@ -956,7 +956,7 @@ better accuracy and result train with more number of epochs.
 test_ds = load(split="sbd_eval")
 test_ds = preprocess_inputs(test_ds)
 
-images, masks = next(iter(train_ds.take(1)))
+images, masks = next(iter(test_ds.take(1)))
 images = ops.convert_to_tensor(images)
 masks = ops.convert_to_tensor(masks)
 preds = ops.expand_dims(ops.argmax(model.predict(images), axis=-1), axis=-1)