mirror of
https://github.com/skindhu/Build-A-Large-Language-Model-CN.git
synced 2026-07-29 15:20:24 +08:00
@@ -1069,6 +1069,7 @@ def generate(model, idx, max_new_tokens, context_size,
|
|||||||
idx_next = torch.argmax(logits, dim=-1, keepdim=True)
|
idx_next = torch.argmax(logits, dim=-1, keepdim=True)
|
||||||
if idx_next == eos_id: #E
|
if idx_next == eos_id: #E
|
||||||
break
|
break
|
||||||
|
idx_next = idx_next.unsqueeze(1)
|
||||||
idx = torch.cat((idx, idx_next), dim=1)
|
idx = torch.cat((idx, idx_next), dim=1)
|
||||||
return idx
|
return idx
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user