Before update the tr_loss, make sure tr_loss_step is in the same devi…

…ce. (#1439) * before update the loss from dpo, make sure it's in the same device of tr_loss * Update trl/trainer/dpo_trainer.py Co-authored-by: guy1992l <[email protected]> --------- Co-authored-by: guy1992l <[email protected]>
huggingface · Mar 19, 2024 · f976c6d · f976c6d
1 parent abc7301
commit f976c6d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
@@ -1080,6 +1080,8 @@ def compute_loss(
         with compute_loss_context_manager():
             loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train")
 
+        # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class:
+        loss = loss.to(self.args.device)
         # force log the metrics
         self.store_metrics(metrics, train_eval="train")