Description
### Summary
`GPT.generate()` divides the logits by `temperature` unconditionally. With
`temperature=0` — the standard convention for deterministic / greedy decoding
(OpenAI, Hugging Face, llama.cpp all treat 0 as "argmax") — this is a division
by zero: the logits become `inf`/`nan`, softmax yields `nan` probabilities, and
`torch.multinomial` raises. So the most natural way to ask for deterministic
output crashes instead of working.
### Steps to reproduce
```python
import torch
from model import GPT, GPTConfig
cfg = GPTConfig(block_size=16, vocab_size=32, n_layer=1, n_head=1,
n_embd=16, dropout=0.0, bias=False)
m = GPT(cfg).eval()
idx = torch.zeros((1, 1), dtype=torch.long)
m.generate(idx, max_new_tokens=5, temperature=0.0) # boomActual behavior
RuntimeError: probability tensor contains either `inf`, `nan` or element < 0
Expected behavior
temperature=0 should perform greedy decoding (pick argmax of the logits),
not crash — consistent with every other major sampling implementation.
Root cause
model.py, in GPT.generate():
logits = logits[:, -1, :] / temperature # temperature == 0 -> inf/nan
Suggested fix
Short-circuit temperature == 0 to greedy decoding before the division:
logits = logits[:, -1, :]
if temperature == 0.0:
idx_next = torch.argmax(logits, dim=-1, keepdim=True)
else:
logits = logits / temperature
# ... existing top-k + multinomial path
Description
Actual behavior
Expected behavior
temperature=0should perform greedy decoding (pickargmaxof the logits),not crash — consistent with every other major sampling implementation.
Root cause
model.py, inGPT.generate():Suggested fix
Short-circuit
temperature == 0to greedy decoding before the division: