fix dpo trainer bug (huggingface#1073)

xzqxnet0990 · Dec 8, 2023 · ac8fee3 · ac8fee3
1 parent 7d0a8ee
commit ac8fee3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
@@ -453,7 +453,7 @@ def collate(self, batch):
                     if "prompt" in k:
                         to_pad = [torch.LongTensor(ex[k][::-1]) for ex in batch]
                     else:
-                        to_pad = [torch.LongTensor(ex[k]) for ex in batch]
+                        to_pad = [torch.LongTensor(ex[k][:-1]) for ex in batch]
                     if k.endswith("_input_ids"):
                         padding_value = self.tokenizer.pad_token_id
                     elif k.endswith("_labels"):