Unverified Commit 3d14707a authored by Stella Biderman's avatar Stella Biderman Committed by GitHub
Browse files

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

hotfix(gpt2): Remove vocab-size logits slice
parents f9eca2c8 1a545447
...@@ -119,7 +119,7 @@ class HFLM(BaseLM): ...@@ -119,7 +119,7 @@ class HFLM(BaseLM):
logits returned from the model logits returned from the model
""" """
with torch.no_grad(): 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): def _model_generate(self, context, max_length, eos_token_id):
return self.gpt2.generate( return self.gpt2.generate(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment