From 18b5724e757e22ce01fd31cf4300e0c1d562740a Mon Sep 17 00:00:00 2001
From: Kostyantyn Borysenko <kostya.borysenko@gmail.com>
Date: Sun, 26 May 2024 13:05:07 -0700
Subject: [PATCH] Fix an incorrect input dimension

---
 ch03/01_main-chapter-code/multihead-attention.ipynb | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ch03/01_main-chapter-code/multihead-attention.ipynb b/ch03/01_main-chapter-code/multihead-attention.ipynb
index 1748607..98eed3f 100644
--- a/ch03/01_main-chapter-code/multihead-attention.ipynb
+++ b/ch03/01_main-chapter-code/multihead-attention.ipynb
@@ -228,7 +228,7 @@
     "            [CausalSelfAttention(d_in, d_out, context_length, dropout, qkv_bias) \n",
     "             for _ in range(num_heads)]\n",
     "        )\n",
-    "        self.out_proj = nn.Linear(d_out*num_heads, d_out*num_heads)\n",
+    "        self.out_proj = nn.Linear(d_in*num_heads, d_out*num_heads)\n",
     "\n",
     "    def forward(self, x):\n",
     "        context_vec = torch.cat([head(x) for head in self.heads], dim=-1)\n",
@@ -383,7 +383,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,