Fix inference (#31)

EleutherAI · Apr 24, 2024 · 80d2998 · 80d2998
1 parent 3091aa9
commit 80d2998
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/amt/inference/model.py b/amt/inference/model.py
@@ -354,6 +354,7 @@ def __init__(
             ]
         )
         self.ln = nn.LayerNorm(n_state)
+        self.output = nn.Linear(n_state, n_vocab, bias=False)
         self.register_buffer("causal_mask", None, persistent=False)
 
     def forward(
@@ -376,9 +377,7 @@ def forward(
             )
 
         x = self.ln(x)
-        logits = (
-            x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)
-        ).float()
+        logits = self.output(x)
 
         return logits