hotfix(gpt2): Remove vocab-size logits slice

1a545447 · jon-tow · 2bae14d2 · 1a545447
Commit 1a545447 authored Feb 04, 2023 by jon-tow
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(