osmarks/nanogpt-experiments
1
0
Fork 0
mirror of https://github.com/osmarks/nanogpt-experiments.git synced 2024-12-18 14:10:28 +00:00
Code Issues Releases Wiki Activity

fix bug... if topk > vocab_size, torch.topk will throw error

Browse Source
This commit is contained in:
Andrej Karpathy 2023-01-14 03:57:00 +00:00
parent 57735f532d
commit 91d02510ce
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
model.py
View File

@ -277,7 +277,7 @@ class GPT(nn.Module):
logits = logits[:, -1, :] / temperature logits = logits[:, -1, :] / temperature
# optionally crop the logits to only the top k options # optionally crop the logits to only the top k options
if top_k is not None: if top_k is not None:
v, _ = torch.topk(logits, top_k) v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf') logits[logits < v[:, [-1]]] = -float('Inf')
# apply softmax to convert logits to (normalized) probabilities # apply softmax to convert logits to (normalized) probabilities
probs = F.softmax(logits, dim=-1) probs = F.softmax(logits, dim=-1)