[MLU] add mlu device for infer

PaddlePaddle · Jul 14, 2023 · e3efae0 · e3efae0
1 parent 991fb7c
commit e3efae0
Showing 1 changed file with 4 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,9 @@ def create_predictor(args, mode, logger):
 
         elif args.use_npu:
             config.enable_custom_device("npu")
+        elif args.use_mlu:
+            dev_type = paddle.device.get_all_custom_device_type()[0]
+            config.enable_custom_device(dev_type)
         elif args.use_xpu:
             config.enable_xpu(10 * 1024 * 1024)
         else: