support mlu for infer (#13617)

PaddlePaddle · Aug 7, 2024 · ac43b3e · ac43b3e
1 parent 3aa1f42
commit ac43b3e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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: