colin
/
Witllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refine.

This commit is contained in:
Colin 2023-12-21 20:52:19 +08:00
parent c462129ba6
commit bfc3fb6706
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
chatglm/modeling_chatglm.py
View File

@ -752,7 +752,7 @@ class ChatGLMForConditionalGeneration(nn.Module):
return_dict=True,
output_hidden_states=output_hidden_states,
)
next_token_logits = logits[:, 0, :]
next_token_logits = logits[:, -1, :]
probs = nn.functional.softmax(next_token_logits, dim=-1)
next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)