From 3cd3473bf7a3b41484baa86d9092248d78e7af39 Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期五, 21 四月 2023 17:17:37 +0800
Subject: [PATCH] docs
---
funasr/train/trainer.py | 10 ++++++++--
1 files changed, 8 insertions(+), 2 deletions(-)
diff --git a/funasr/train/trainer.py b/funasr/train/trainer.py
index 9c4af41..7c187e9 100644
--- a/funasr/train/trainer.py
+++ b/funasr/train/trainer.py
@@ -582,10 +582,16 @@
if num_batch_updates % batch_interval == 0:
if options.use_pai and options.oss_bucket is not None:
buffer = BytesIO()
- torch.save(model.state_dict(), buffer)
+ if hasattr(model, "module"):
+ torch.save(model.module.state_dict(), buffer)
+ else:
+ torch.save(model.state_dict(), buffer)
options.oss_bucket.put_object(os.path.join(output_dir, f"{num_batch_updates}step.pb"), buffer.getvalue())
else:
- torch.save(model.state_dict(), os.path.join(output_dir, f"{num_batch_updates}step.pb"))
+ if hasattr(model, "module"):
+ torch.save(model.module.state_dict(), os.path.join(output_dir, f"{num_batch_updates}step.pb"))
+ else:
+ torch.save(model.state_dict(), os.path.join(output_dir, f"{num_batch_updates}step.pb"))
if distributed:
torch.distributed.all_reduce(iterator_stop, ReduceOp.SUM)
--
Gitblit v1.9.1