diff --git a/ch04/01_main-chapter-code/ch04.ipynb b/ch04/01_main-chapter-code/ch04.ipynb
index 81f931d..5001926 100644
--- a/ch04/01_main-chapter-code/ch04.ipynb
+++ b/ch04/01_main-chapter-code/ch04.ipynb
@@ -1235,7 +1235,7 @@
     "        # Apply softmax to get probabilities\n",
     "        probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)\n",
     "\n",
-    "        # Get the idx of the vocab entry with the highest logits value\n",
+    "        # Get the idx of the vocab entry with the highest probability value\n",
     "        idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # (batch, 1)\n",
     "\n",
     "        # Append sampled index to the running sequence\n",