move empty cache

2025-06-24 06:29:09 +00:00 · 2025-04-10 10:22:43 +08:00 · 2025-04-10 10:22:43 +08:00 · 1723a02860
commit 1723a02860
parent f7e532511c
1 changed files with 1 additions and 0 deletions
--- a/applications/ColossalChat/coati/distributed/consumer.py
+++ b/applications/ColossalChat/coati/distributed/consumer.py
@ -129,6 +129,7 @@ class BaseConsumer:

                    if episode != self.num_episodes - 1 or step != self.num_update_per_episode - 1:
                        print(f"[T{dist.get_rank()}] Sync model episode {episode} step {step}")
+                        torch.cuda.empty_cache()
                        state_dict = self.state_dict()
                        if self.rank == 0:
                            ray_broadcast_tensor_dict(