fix step in continue training in steps mode (#2453)

wenet-e2e · Mar 29, 2024 · f3a434a · f3a434a
1 parent 377e720
commit f3a434a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/wenet/utils/train_utils.py b/wenet/utils/train_utils.py
@@ -605,7 +605,9 @@ def log_per_step(writer, info_dict, timer: Optional[StepTimer] = None):
             steps_per_second = timer.steps_per_second(timer_step)
             log_str += 'steps/sec {:.1f}| '.format(steps_per_second)
         log_str += 'Batch {}/{} loss {:.6f} '.format(
-            epoch, batch_idx + 1, loss_dict['loss'] * accum_grad)
+            epoch,
+            batch_idx + 1 if 'save_interval' not in info_dict else step + 1,
+            loss_dict['loss'] * accum_grad)
         for name, value in loss_dict.items():
             if name != 'loss' and value is not None:
                 log_str += '{} {:.6f} '.format(name, value)