[CPO] fix memory leak due to retained value (huggingface#1531)

lapp0 · May 10, 2024 · a966d01 · a966d01
1 parent abf18db
commit a966d01
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trl/trainer/cpo_trainer.py b/trl/trainer/cpo_trainer.py
@@ -738,7 +738,7 @@ def get_batch_loss_metrics(
         metrics[f"{prefix}logps/chosen"] = policy_chosen_logps.detach().mean().cpu()
         metrics[f"{prefix}logits/rejected"] = policy_rejected_logits.detach().mean().cpu()
         metrics[f"{prefix}logits/chosen"] = policy_chosen_logits.detach().mean().cpu()
-        metrics[f"{prefix}nll_loss"] = policy_nll_loss.cpu().mean()
+        metrics[f"{prefix}nll_loss"] = policy_nll_loss.detach().mean().cpu()
 
         return loss, metrics