Update train.py

train.py

@@ -68,7 +68,7 @@ def preprocess_data(example):
 # 预处理数据集
 dataset = dataset.map(preprocess_data, batched=True)
 
-# ✅ 自定义 `DistillationTrainer`，覆盖 `training_step()` 以防止 `num_items_in_batch` 传递
+# ✅ 修正 training_step() 参数问题
 class DistillationTrainer(Trainer):
     def __init__(self, teacher, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -104,7 +104,7 @@ class DistillationTrainer(Trainer):
 
         return (loss, outputs_student) if return_outputs else loss
 
-    def training_step(self, model, inputs):
+    def training_step(self, model, inputs, *args, **kwargs):  # ✅ 修正：添加 *args, **kwargs 以兼容 Trainer
         """✅ 关键修复点：覆盖 `training_step()`，防止 `num_items_in_batch` 传递"""
         model.train()
         inputs = self._prepare_inputs(inputs)