From f7e532511ce766bb1f5661c397fb468c3bad5523 Mon Sep 17 00:00:00 2001 From: YeAnbang Date: Thu, 10 Apr 2025 10:22:28 +0800 Subject: [PATCH] move empty cache --- applications/ColossalChat/coati/distributed/grpo_consumer.py | 1 - 1 file changed, 1 deletion(-) diff --git a/applications/ColossalChat/coati/distributed/grpo_consumer.py b/applications/ColossalChat/coati/distributed/grpo_consumer.py index 68a01528e..e23254d1b 100644 --- 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(