diff --git a/pytorch_unet_resnet18_colab.ipynb b/pytorch_unet_resnet18_colab.ipynb
index fc64cc4..b2f2cdf 100644
--- a/pytorch_unet_resnet18_colab.ipynb
+++ b/pytorch_unet_resnet18_colab.ipynb
@@ -424,7 +424,7 @@
         "import simulation\n",
         "\n",
         "# Generate some random images\n",
-        "input_images, target_masks = simulation.generate_random_data(192, 192, count=3)\n",
+        "input_images, target_masks = simulation.generate_random_data(32, 32, count=3)\n",
         "\n",
         "print(\"input_images shape and range\", input_images.shape, input_images.min(), input_images.max())\n",
         "print(\"target_masks shape and range\", target_masks.shape, target_masks.min(), target_masks.max())\n",
@@ -511,7 +511,7 @@
         "\n",
         "class SimDataset(Dataset):\n",
         "  def __init__(self, count, transform=None):\n",
-        "    self.input_images, self.target_masks = simulation.generate_random_data(192, 192, count=count)\n",
+        "    self.input_images, self.target_masks = simulation.generate_random_data(32, 32, count=count)\n",
         "    self.transform = transform\n",
         "\n",
         "  def __len__(self):\n",
@@ -1645,4 +1645,4 @@
       "outputs": []
     }
   ]
-}
\ No newline at end of file
+}