From f690741b9aa31a27d6620b02eb1c754d6fe27237 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期三, 01 五月 2024 01:05:59 +0800
Subject: [PATCH] log step

---
 funasr/train_utils/trainer.py |    6 ++++--
 1 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/funasr/train_utils/trainer.py b/funasr/train_utils/trainer.py
index e0b6def..01e2924 100644
--- a/funasr/train_utils/trainer.py
+++ b/funasr/train_utils/trainer.py
@@ -465,7 +465,8 @@
                     batch_num_epoch = len(dataloader_train)
                 self.log(
                     epoch,
-                    batch_idx + kwargs.get("start_step", 0),
+                    batch_idx,
+                    log_step=batch_idx + kwargs.get("start_step", 0),
                     step_in_epoch=self.step_in_epoch,
                     batch_num_epoch=batch_num_epoch,
                     lr=lr,
@@ -634,11 +635,12 @@
         tag="train",
         data_split_i=0,
         data_split_num=1,
+        log_step=None,
         **kwargs,
     ):
 
         if (batch_idx + 1) % self.log_interval == 0:
-
+            batch_idx = log_step if log_step is not None else batch_idx
             gpu_info = (
                 "GPU, memory: usage: {:.3f} GB, "
                 "peak: {:.3f} GB, "

--
Gitblit v1.9.1