From d097d0ca45472965d4411357d52adda5657691a2 Mon Sep 17 00:00:00 2001
From: R1ckShi <shixian.shi@alibaba-inc.com>
Date: 星期四, 30 五月 2024 14:59:07 +0800
Subject: [PATCH] update
---
funasr/utils/export_utils.py | 7 +++++--
1 files changed, 5 insertions(+), 2 deletions(-)
diff --git a/funasr/utils/export_utils.py b/funasr/utils/export_utils.py
index 7d6606b..ba200a6 100644
--- a/funasr/utils/export_utils.py
+++ b/funasr/utils/export_utils.py
@@ -78,12 +78,15 @@
)
-def _torchscripts(model, path, device='cpu'):
+def _torchscripts(model, path, device='cuda'):
dummy_input = model.export_dummy_inputs()
if device == 'cuda':
model = model.cuda()
- dummy_input = tuple([i.cuda() for i in dummy_input])
+ if isinstance(dummy_input, torch.Tensor):
+ dummy_input = dummy_input.cuda()
+ else:
+ dummy_input = tuple([i.cuda() for i in dummy_input])
# model_script = torch.jit.script(model)
model_script = torch.jit.trace(model, dummy_input)
--
Gitblit v1.9.1