Update trl/trainer/dpo_trainer.py

huggingface · Jan 7, 2025 · 9f086e5 · 9f086e5
1 parent 175c5e2
commit 9f086e5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
@@ -1170,7 +1170,8 @@ def concatenated_forward(self, model: nn.Module, batch: dict[str, Union[list, to
                 #              [d, e, f, g]]     position_ids = [[0, 1, 2, 0, 1, 2, 3]]
                 input_ids = input_ids[attention_mask.bool()].unsqueeze(0)
                 loss_mask = loss_mask[attention_mask.bool()].unsqueeze(0)
-                model_kwargs["position_ids"] = attention_mask.cumsum(1)[attention_mask.bool()].unsqueeze(0) - 1
+                position_ids = attention_mask.cumsum(1)[attention_mask.bool()].unsqueeze(0) - 1
+                model_kwargs["position_ids"] = position_ids
             else:
                 model_kwargs["attention_mask"] = attention_mask