diff --git a/ch03/01_main-chapter-code/multihead-attention.ipynb b/ch03/01_main-chapter-code/multihead-attention.ipynb
index fbbdee8..6992fce 100644
--- a/ch03/01_main-chapter-code/multihead-attention.ipynb
+++ b/ch03/01_main-chapter-code/multihead-attention.ipynb
@@ -88,7 +88,7 @@
     "\n",
     "\n",
     "token_embedding_layer = nn.Embedding(vocab_size, output_dim)\n",
-    "pos_embedding_layer = nn.Embedding(vocab_size, output_dim)\n",
+    "pos_embedding_layer = torch.nn.Embedding(block_size, output_dim)\n",
     "\n",
     "max_length = 4\n",
     "dataloader = create_dataloader(raw_text, batch_size=8, max_length=max_length, stride=5)"
@@ -324,6 +324,14 @@
     "\n",
     "print(\"context_vecs.shape:\", context_vecs.shape)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f8d4be84-28bb-41d5-996c-4936acffd411",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {