diff --git a/ch04/01_main-chapter-code/ch04.ipynb b/ch04/01_main-chapter-code/ch04.ipynb
index 2eaecee..7111d55 100644
--- a/ch04/01_main-chapter-code/ch04.ipynb
+++ b/ch04/01_main-chapter-code/ch04.ipynb
@@ -756,7 +756,7 @@
     "            # Compute the output of the current layer\n",
     "            layer_output = layer(x)\n",
     "            # Check if shortcut can be applied\n",
-    "            if self.use_shortcut and x.size() == layer_output.size():\n",
+    "            if self.use_shortcut and x.shape == layer_output.shape:\n",
     "                x = x + layer_output\n",
     "            else:\n",
     "                x = layer_output\n",