move empty cache

2025-07-31 23:35:47 +00:00 · 2025-04-10 10:22:28 +08:00 · 2025-04-10 10:22:28 +08:00 · f7e532511c
commit f7e532511c
parent ff3ca30a35
1 changed files with 0 additions and 1 deletions
--- a/applications/ColossalChat/coati/distributed/grpo_consumer.py
+++ b/applications/ColossalChat/coati/distributed/grpo_consumer.py
@ -285,7 +285,6 @@ class GRPOConsumer(BaseConsumer):
                            kl = all_reduce_mean(torch.mean(torch.stack(kl)).to(loss.device), self.plugin).data
                            mean_kl.append(kl)
                        mean_loss.append(all_reduce_mean(loss, self.plugin).data)
-                    torch.cuda.empty_cache()
                else:

                    policy_model_logits = self.policy_model(