From e41a6b7c28bbbf9ebbbde79e283465d7b266f3be Mon Sep 17 00:00:00 2001 From: "877825076@qq.com" <877825076@qq.com> Date: Wed, 10 Jan 2024 15:50:17 +0800 Subject: [PATCH] fix --- internlm/utils/model_checkpoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internlm/utils/model_checkpoint.py b/internlm/utils/model_checkpoint.py index 4186ef6..2da2755 100644 --- a/internlm/utils/model_checkpoint.py +++ b/internlm/utils/model_checkpoint.py @@ -996,7 +996,7 @@ now step_count is {train_state.step_count}", if ( train_state.step_count > 0 and train_state.step_count % self.checkpoint_every == 0 - or train_state.step_count == train_state.total_steps + or train_state.batch_count == (train_state.total_steps - 1) ): save_ckpts, save_type = True, CheckpointSaveType.NORMAL_CHECKPOINT