Committed by
GitHub
Fix: export-onnx.py(expected all tensors to be on the same device) (#1699)
由于SenseVoiceSmall.from_pretrained()
调用的funasr.auto.auto_model.AutoModel.build_model()默认device是cuda
(在cuda available的环境中)
```py
device = kwargs.get("device", "cuda")
if not torch.cuda.is_available() or kwargs.get("ngpu", 1) == 0:
device = "cpu"
kwargs["batch_size"] = 1
kwargs["device"] = device
```
而export-onnx.py里的tensor默认都是cpu, 导致
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu
所以直接在加载model的时候指定cpu
正在显示
1 个修改的文件
包含
1 行增加
和
1 行删除
| @@ -119,7 +119,7 @@ def display_params(params): | @@ -119,7 +119,7 @@ def display_params(params): | ||
| 119 | 119 | ||
| 120 | 120 | ||
| 121 | def main(): | 121 | def main(): |
| 122 | - model, params = SenseVoiceSmall.from_pretrained(model="iic/SenseVoiceSmall") | 122 | + model, params = SenseVoiceSmall.from_pretrained(model="iic/SenseVoiceSmall", device="cpu") |
| 123 | display_params(params) | 123 | display_params(params) |
| 124 | 124 | ||
| 125 | generate_tokens(params) | 125 | generate_tokens(params) |
-
请 注册 或 登录 后发表评论