tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-07-09 22:39:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix generate() crash when top_k=0

Browse Source 
Prevent a crash in generate() by skipping top-k filtering when top_k is set to 0
This commit is contained in:
Harsh Gupta 2026-01-28 14:31:30 +05:30 committed by GitHub
parent c88bbf8133
commit d4d014010b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nanochat/gpt.py
View File

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