Skip to content

GitLab

Commit f298ca76 authored by Leo Gao's avatar Leo Gao
Browse files

Fix memory problem

parent 984f617b
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
lm_eval/models/gpt2.py
View file @ f298ca76
...@@ -32,7 +32,7 @@ class GPT2LM(LM): ...@@ -32,7 +32,7 @@ class GPT2LM(LM):
logits = F.log_softmax(self.gpt2(inp)[0], dim=-1)[:, ctxlen - 1:-1] # [batch, seq, vocab] logits = F.log_softmax(self.gpt2(inp)[0], dim=-1)[:, ctxlen - 1:-1] # [batch, seq, vocab]
# TODO: implement isgreedy # TODO: implement isgreedy
res.append((torch.gather(logits, 2, cont_toks.unsqueeze(-1)).squeeze(-1), False)) res.append((float(torch.gather(logits, 2, cont_toks.unsqueeze(-1)).squeeze(-1)), False))
return res return res
......
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