tiny fix

zhanglv0209 · May 8, 2023 · 995c032 · 995c032
1 parent b9b0353
commit 995c032
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/train_rm.py b/src/train_rm.py
@@ -27,6 +27,8 @@ def main():
         inference_mode=(not training_args.do_train)
     )
 
+    training_args.remove_unused_columns = False # Important for pairwise dataset
+
     # Initialize our Trainer
     trainer = PairwiseTrainerForChatGLM(
         finetuning_args=finetuning_args,