Update dpo_trainer.py (#1049)

huggingface · Dec 1, 2023 · a60ceef · a60ceef
1 parent baa8f09
commit a60ceef
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
@@ -632,15 +632,15 @@ def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[
         if self.ref_model is None:
             with self.accelerator.unwrap_model(self.model).disable_adapter():
                 reference_output = self.model.generate(
-                    batch["prompt_input_ids"],
+                    input_ids=batch["prompt_input_ids"],
                     attention_mask=batch["prompt_attention_mask"],
                     max_length=self.max_length,
                     do_sample=True,
                     pad_token_id=self.tokenizer.pad_token_id,
                 )
         else:
             reference_output = self.ref_model.generate(
-                batch["prompt_input_ids"],
+                input_ids=batch["prompt_input_ids"],
                 attention_mask=batch["prompt_attention_mask"],
                 max_length=self.max_length,
                 do_sample=True,