From 758696c59cc904585c05be8bafd6296a9c10fe04 Mon Sep 17 00:00:00 2001 From: tsukumi Date: Mon, 23 Sep 2024 12:40:07 +0900 Subject: [PATCH] Fix: Device ID specified in ExecutionProvider options is respected when moving input tensors --- style_bert_vits2/models/infer_onnx.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/style_bert_vits2/models/infer_onnx.py b/style_bert_vits2/models/infer_onnx.py index e147395..259a51d 100644 --- a/style_bert_vits2/models/infer_onnx.py +++ b/style_bert_vits2/models/infer_onnx.py @@ -195,10 +195,17 @@ def infer_onnx( else: 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 メモリに入力テンソルを割り当て io_binding = onnx_session.io_binding() 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) # 推論の実行