Merge pull request #384 from jon-tow/gpt2-vocab-slice

hotfix(gpt2): Remove vocab-size logits slice

Merge pull request #384 from jon-tow/gpt2-vocab-slice
hotfix(gpt2): Remove vocab-size logits slice
3d14707a · Stella Biderman · GitHub · f9eca2c8 · 1a545447 · 3d14707a
Unverified Commit 3d14707a authored Feb 05, 2023 by Stella Biderman Committed by GitHub Feb 05, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

lm_eval/models/gpt2.py lm_eval/models/gpt2.py +1 -1

No files found.
--- a/lm_eval/models/gpt2.py
+++ b/lm_eval/models/gpt2.py
@@ -119,7 +119,7 @@ class HFLM(BaseLM):
        logits returned from the model
        """
        with torch.no_grad():
-            return self.gpt2(inps)[0][:, :, :50257]
+            return self.gpt2(inps)[0]

    def _model_generate(self, context, max_length, eos_token_id):
        return self.gpt2.generate(