From a2d17671ec4de579de861b41d80340c4d7c364b6 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Tue, 22 Apr 2025 14:11:15 +0800 Subject: [PATCH] add final save --- applications/ColossalChat/coati/distributed/consumer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/applications/ColossalChat/coati/distributed/consumer.py b/applications/ColossalChat/coati/distributed/consumer.py index 79beb2a2d..b7b865b26 100644 --- a/applications/ColossalChat/coati/distributed/consumer.py +++ b/applications/ColossalChat/coati/distributed/consumer.py @@ -119,7 +119,7 @@ class BaseConsumer: assert len(self.buffer) == 0 if self.lr_scheduler is not None: self.lr_scheduler.step() - if (step + 1) % self.save_interval == 0: + if (step + 1) % self.save_interval == 0 or (step + 1) == self.num_update_per_episode: if self.rank == 0: print(f"Start saving policy model at step {step + 1}.") save_path = os.path.join(self.save_dir, f"modeling-step-{step + 1}")