tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

flipped the slice

Browse Source
This commit is contained in:
Wing Lian
2024-12-19 01:21:48 -05:00
parent 283faf3909
commit 1c603da96a
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/axolotl/core/trainers/kd.py
View File

@@ -23,7 +23,7 @@ def kd_loss_function(
# Slice student logits to match the teacher-provided sequence length # Slice student logits to match the teacher-provided sequence length
student_logits_for_kd = student_logits[ student_logits_for_kd = student_logits[
:, -teacher_seq_len:, : :, :teacher_seq_len, :
] # [B, teacher_seq_len, vocab_size] ] # [B, teacher_seq_len, vocab_size]
# Gather student logits for teacher's top-K tokens # Gather student logits for teacher's top-K tokens