From 67c1c8405e12a8f909751df4021e2a5da684f927 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 21 Jan 2025 11:23:38 -0500 Subject: [PATCH] use iter instead of tuple --- src/axolotl/integrations/kd/chat_template.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/axolotl/integrations/kd/chat_template.py b/src/axolotl/integrations/kd/chat_template.py index 5a0c4e90b..2f17f0a9b 100644 --- a/src/axolotl/integrations/kd/chat_template.py +++ b/src/axolotl/integrations/kd/chat_template.py @@ -79,7 +79,7 @@ class ChatTemplateStrategyWithKD(ChatTemplateStrategy): # target_seq_len = input_seq_len # truncate the second dimension of the logprobs to top_k - logprobs = logprobs[:, :top_k] + logprobs = [row[:top_k] for row in logprobs] # fill with -inf for padding_len tokens for top_k tokens # extend target_logprobs with a padding_len x top_k 2D list filled with -inf