From bfc3fb67069c512b82cacdf6e0a2b1a107876b32 Mon Sep 17 00:00:00 2001
From: Colin <colin>
Date: Thu, 21 Dec 2023 20:52:19 +0800
Subject: [PATCH] Refine.

---
 chatglm/modeling_chatglm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/chatglm/modeling_chatglm.py b/chatglm/modeling_chatglm.py
index 292861d..653f2f1 100644
--- a/chatglm/modeling_chatglm.py
+++ b/chatglm/modeling_chatglm.py
@@ -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)