tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Fix: Device ID specified in ExecutionProvider options is respected when moving input tensors

Browse Source
This commit is contained in:
tsukumi
2024-09-23 12:40:07 +09:00
parent c792bf9a0d
commit 758696c59c
1 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
style_bert_vits2/models/infer_onnx.py
View File

@@ -195,10 +195,17 @@ def infer_onnx(
else: else:
device_type = "cpu" device_type = "cpu"
device_id = 0
first_provider_options = onnx_session.get_provider_options()[first_provider]
if "device_id" in first_provider_options:
device_id = int(first_provider_options["device_id"])
# GPU メモリに入力テンソルを割り当て # GPU メモリに入力テンソルを割り当て
io_binding = onnx_session.io_binding() io_binding = onnx_session.io_binding()
for name, value in zip(input_names, input_tensor): for name, value in zip(input_names, input_tensor):
gpu_tensor = onnxruntime.OrtValue.ortvalue_from_numpy(value, device_type) gpu_tensor = onnxruntime.OrtValue.ortvalue_from_numpy(
value, device_type, device_id
)
io_binding.bind_ortvalue_input(name, gpu_tensor) io_binding.bind_ortvalue_input(name, gpu_tensor)
# 推論の実行 # 推論の実行