From ac43b3eb6195ea795b350d9490226c48dc2308f7 Mon Sep 17 00:00:00 2001 From: wangna11BD <79366697+wangna11BD@users.noreply.github.com> Date: Wed, 7 Aug 2024 16:46:13 +0800 Subject: [PATCH] support mlu for infer (#13617) --- tools/infer/utility.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tools/infer/utility.py b/tools/infer/utility.py index fd3567730b..4d34680317 100644 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -37,6 +37,7 @@ def init_args(): parser.add_argument("--use_gpu", type=str2bool, default=True) parser.add_argument("--use_xpu", type=str2bool, default=False) parser.add_argument("--use_npu", type=str2bool, default=False) + parser.add_argument("--use_mlu", type=str2bool, default=False) parser.add_argument("--ir_optim", type=str2bool, default=True) parser.add_argument("--use_tensorrt", type=str2bool, default=False) parser.add_argument("--min_subgraph_size", type=int, default=15) @@ -247,6 +248,8 @@ def create_predictor(args, mode, logger): elif args.use_npu: config.enable_custom_device("npu") + elif args.use_mlu: + config.enable_custom_device("mlu") elif args.use_xpu: config.enable_xpu(10 * 1024 * 1024) else: