From 5bd8cd7aed013fe4893e118677d884c28d4a342a Mon Sep 17 00:00:00 2001 From: Zhikang Niu <73390819+ZhikangNiu@users.noreply.github.com> Date: Fri, 28 Mar 2025 13:53:12 +0800 Subject: [PATCH] update: better save last & per ckpt logic #924 Co-authored-by: Yushen CHEN <45333109+SWivid@users.noreply.github.com> --- src/f5_tts/model/trainer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/f5_tts/model/trainer.py b/src/f5_tts/model/trainer.py index 7a405c2..5f97d30 100644 --- a/src/f5_tts/model/trainer.py +++ b/src/f5_tts/model/trainer.py @@ -395,6 +395,9 @@ class Trainer: self.writer.add_scalar("loss", loss.item(), global_update) self.writer.add_scalar("lr", self.scheduler.get_last_lr()[0], global_update) + if global_update % self.last_per_updates == 0 and self.accelerator.sync_gradients: + self.save_checkpoint(global_update, last=True) + if global_update % self.save_per_updates == 0 and self.accelerator.sync_gradients: self.save_checkpoint(global_update) @@ -430,9 +433,6 @@ class Trainer: ) self.model.train() - if global_update % self.last_per_updates == 0 and self.accelerator.sync_gradients: - self.save_checkpoint(global_update, last=True) - self.save_checkpoint(global_update, last=True) self.accelerator.end_training()