From cb8b09e085bdfb5599ccd1a862912bc0b7e4f41c Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期二, 20 二月 2024 16:42:09 +0800
Subject: [PATCH] update

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

diff --git a/funasr/train_utils/trainer.py b/funasr/train_utils/trainer.py
index 3cd61a1..cf997a4 100644
--- a/funasr/train_utils/trainer.py
+++ b/funasr/train_utils/trainer.py
@@ -273,8 +273,9 @@
                 speed_stats["total_time"] = total_time
 
 
-            pbar.update(1)
+            
             if self.local_rank == 0:
+                pbar.update(1)
                 gpu_info = "GPU, memory: {:.3f} GB, " \
                            "{:.3f} GB, "\
                            "{:.3f} GB, "\
@@ -290,6 +291,7 @@
                     f"(loss: {loss.detach().cpu().item():.3f}), "
                     f"{[(k, round(v.cpu().item(), 3)) for k, v in stats.items()]}"
                     f"{gpu_info}"
+                    f"rank: {self.local_rank}"
                 )
                 pbar.set_description(description)
                 if self.writer:
@@ -344,14 +346,16 @@
                 loss = loss
                 time4 = time.perf_counter()
 
-                pbar.update(1)
+                
                 if self.local_rank == 0:
+                    pbar.update(1)
                     description = (
                         f"validation epoch: {epoch}/{self.max_epoch}, "
                         f"step {batch_idx}/{len(self.dataloader_train)}, "
                         f"{speed_stats}, "
                         f"(loss: {loss.detach().cpu().item():.3f}), "
                         f"{[(k, round(v.cpu().item(), 3)) for k, v in stats.items()]}"
+                        f"rank: {self.local_rank}"
                     )
                     pbar.set_description(description)
                     if self.writer:

--
Gitblit v1.9.1