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 | 5 ++++-
1 files changed, 4 insertions(+), 1 deletions(-)
diff --git a/funasr/utils/export_utils.py b/funasr/utils/export_utils.py
index 1deeaa7..ba200a6 100644
--- a/funasr/utils/export_utils.py
+++ b/funasr/utils/export_utils.py
@@ -83,7 +83,10 @@
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