Commit 1a545447 authored by jon-tow's avatar jon-tow
Browse files

hotfix(gpt2): Remove vocab-size logits slice

parent 2bae14d2
...@@ -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