From 54c3c94b26b02fe878ff6ffb7096ba1c3d22d72a Mon Sep 17 00:00:00 2001
From: speech_asr <wangjiaming.wjm@alibaba-inc.com>
Date: 星期四, 20 四月 2023 17:55:57 +0800
Subject: [PATCH] update
---
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