karpathy#71 use config n_head instead of hardcoded 4 heads

SpeedCoder5 · Apr 25, 2022 · 0ac226d · 0ac226d
1 parent 3ed14b2
commit 0ac226d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mingpt/model.py b/mingpt/model.py
@@ -87,9 +87,9 @@ def __init__(self, config):
         self.ln2 = nn.LayerNorm(config.n_embd)
         self.attn = CausalSelfAttention(config)
         self.mlp = nn.Sequential(
-            nn.Linear(config.n_embd, 4 * config.n_embd),
+            nn.Linear(config.n_embd, config.n_head * config.n_embd),
             nn.GELU(),
-            nn.Linear(4 * config.n_embd, config.n_embd),
+            nn.Linear(config.n_head * config.n_embd, config.n_embd),
             nn.Dropout(config.resid_pdrop),
         )