From e65b1f701abca03bf3a1b5fbb200392aabd38c22 Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期四, 20 六月 2024 17:09:33 +0800
Subject: [PATCH] Dev gzf deepspeed (#1833)
---
funasr/utils/export_utils.py | 5 +++--
1 files changed, 3 insertions(+), 2 deletions(-)
diff --git a/funasr/utils/export_utils.py b/funasr/utils/export_utils.py
index 6583325..5c2a9f4 100644
--- a/funasr/utils/export_utils.py
+++ b/funasr/utils/export_utils.py
@@ -5,7 +5,7 @@
try:
import torch_blade
except Exception as e:
- print(f"failed to load torch_blade: {e}")
+ print(f"Warning, if you are exporting bladedisc, please install it and try it again: pip install -U torch_blade\n")
def export(model, data_in=None, quantize: bool = False, opset_version: int = 14, type='onnx', **kwargs):
@@ -28,6 +28,7 @@
)
elif type == 'torchscripts':
device = 'cuda' if torch.cuda.is_available() else 'cpu'
+ print("Exporting torchscripts on device {}".format(device))
_torchscripts(
m,
path=export_dir,
@@ -195,4 +196,4 @@
model.encoder = _bladedisc_opt(model.encoder, input_data[:2])
model.decoder = _bladedisc_opt(model.decoder, tuple(decoder_inputs))
model_script = torch.jit.trace(model, input_data)
- model_script.save(os.path.join(path, f"{model.export_name}.torchscripts"))
+ model_script.save(os.path.join(path, f"{model.export_name}_blade.torchscripts"))
\ No newline at end of file
--
Gitblit v1.9.1