From 3d5ed00c60e367ced7d0f79654f4c0316326e157 Mon Sep 17 00:00:00 2001 From: lym0302 Date: Mon, 11 Jul 2022 02:53:55 +0000 Subject: [PATCH] specify id, test=doc --- paddlespeech/server/utils/onnx_infer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/paddlespeech/server/utils/onnx_infer.py b/paddlespeech/server/utils/onnx_infer.py index 23d83c73..25802f62 100644 --- a/paddlespeech/server/utils/onnx_infer.py +++ b/paddlespeech/server/utils/onnx_infer.py @@ -30,7 +30,9 @@ def get_sess(model_path: Optional[os.PathLike]=None, sess_conf: dict=None): # "gpu:0" providers = ['CPUExecutionProvider'] if "gpu" in sess_conf.get("device", ""): - providers = ['CUDAExecutionProvider'] + device_id = int(sess_conf["device"].split(":")[1]) + providers = [('CUDAExecutionProvider', {'device_id': device_id})] + # fastspeech2/mb_melgan can't use trt now! if sess_conf.get("use_trt", 0): providers = ['TensorrtExecutionProvider'] -- GitLab