From e98e10639d90c55a4b7e498d0d87837ad9c4173d Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期三, 06 十二月 2023 19:42:02 +0800
Subject: [PATCH] funasr2
---
funasr/cli/trainer.py | 9 +++++----
1 files changed, 5 insertions(+), 4 deletions(-)
diff --git a/funasr/cli/trainer.py b/funasr/cli/trainer.py
index 30e0419..74e058f 100644
--- a/funasr/cli/trainer.py
+++ b/funasr/cli/trainer.py
@@ -131,7 +131,7 @@
for batch_idx, batch in enumerate(self.dataloader_train):
batch = to_device(batch, self.device)
- my_context = model.no_sync if batch_idx % accumulation_steps != 0 else nullcontext
+ my_context = self.model.no_sync if batch_idx % accumulation_steps != 0 else nullcontext
with my_context():
retval = self.model(**batch)
loss, stats, weight = retval
@@ -163,9 +163,10 @@
self.optim.zero_grad()
pbar.update(1)
- pbar.set_description(
- f"Training Epoch: {epoch + 1}/{self.max_epoch}, step {batch_idx}/{len(self.dataloader_train)} (loss: {loss.detach().float()})")
-
+ if self.local_rank == 0:
+ pbar.set_description(
+ f"Training Epoch: {epoch + 1}/{self.max_epoch}, step {batch_idx}/{len(self.dataloader_train)} (loss: {loss.detach().float()})")
+
pbar.close()
# def _train_epoch(self, epoch):
--
Gitblit v1.9.1