This commit is contained in:
Harsh Gupta 2026-01-28 12:11:57 -08:00 committed by GitHub
commit a9db419e90
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -437,7 +437,7 @@ class GPT(nn.Module):
for _ in range(max_tokens):
logits = self.forward(ids) # (B, T, vocab_size)
logits = logits[:, -1, :] # (B, vocab_size)
if top_k is not None:
if top_k is not None and top_k > 0:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')
if temperature > 0: