diff --git a/train_shapes.ipynb b/train_shapes.ipynb
index 82b643c9ec8d7ade40c238f774d83a145c0b1520..25ff3c7ea46a4902939b82400e80448c5ab563e5 100644
--- a/train_shapes.ipynb
+++ b/train_shapes.ipynb
@@ -280,15 +280,15 @@
     "        # Get the center x, y and the size s\n",
     "        x, y, s = dims\n",
     "        if shape == 'square':\n",
-    "            image = cv2.rectangle(image, (x-s, y-s), (x+s, y+s), color, -1)\n",
+    "            cv2.rectangle(image, (x-s, y-s), (x+s, y+s), color, -1)\n",
     "        elif shape == \"circle\":\n",
-    "            image = cv2.circle(image, (x, y), s, color, -1)\n",
+    "            cv2.circle(image, (x, y), s, color, -1)\n",
     "        elif shape == \"triangle\":\n",
     "            points = np.array([[(x, y-s),\n",
     "                                (x-s/math.sin(math.radians(60)), y+s),\n",
     "                                (x+s/math.sin(math.radians(60)), y+s),\n",
     "                                ]], dtype=np.int32)\n",
-    "            image = cv2.fillPoly(image, points, color)\n",
+    "            cv2.fillPoly(image, points, color)\n",
     "        return image\n",
     "\n",
     "    def random_shape(self, height, width):\n",