force return_dict

huggingface · Jan 9, 2024 · c9ed225 · c9ed225
1 parent 989f4a4
commit c9ed225
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py
@@ -220,10 +220,12 @@ def compute_loss(
         rewards_chosen = model(
             input_ids=inputs["input_ids_chosen"],
             attention_mask=inputs["attention_mask_chosen"],
+            return_dict=True,
         )["logits"]
         rewards_rejected = model(
             input_ids=inputs["input_ids_rejected"],
             attention_mask=inputs["attention_mask_rejected"],
+            return_dict=True,
         )["logits"]
         # calculate loss, optionally modulate with margin
         if "margin" in inputs: