diff --git a/ch03/01_main-chapter-code/ch03.ipynb b/ch03/01_main-chapter-code/ch03.ipynb
index 1ba30a9..12019b5 100644
--- a/ch03/01_main-chapter-code/ch03.ipynb
+++ b/ch03/01_main-chapter-code/ch03.ipynb
@@ -1369,7 +1369,7 @@
     }
    ],
    "source": [
-    "row_sums = masked_simple.sum(dim=1, keepdim=True)\n",
+    "row_sums = masked_simple.sum(dim=-1, keepdim=True)\n",
     "masked_simple_norm = masked_simple / row_sums\n",
     "print(masked_simple_norm)"
    ]
@@ -2027,7 +2027,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.6"
+   "version": "3.11.4"
   }
  },
  "nbformat": 4,