fix pp memory issue (#6344)

Co-authored-by: Tong Li <tong.li35271158@gmail.com>
2025-09-09 13:00:52 +00:00 · 2025-06-11 17:54:18 +08:00
parent bb6f5d98fc
commit 8992def757
1 changed files with 1 additions and 1 deletions
--- a/applications/ColossalChat/coati/distributed/grpo_consumer.py
+++ b/applications/ColossalChat/coati/distributed/grpo_consumer.py
@@ -345,7 +345,7 @@ class GRPOConsumer(BaseConsumer):
                        criterion=_criterion,
                        optimizer=self.optimizer,
                        return_loss=True,
-                        return_outputs=True,
+                        return_outputs=False,
                    )
                    loss = policy_model_outputs["loss"]