fix npu inference

pull/3804/head
warrentdrew 3 months ago
parent 7b5218edf2
commit 13ea0dae56

@ -112,7 +112,7 @@ def parse_args():
parser.add_argument(
"--device",
default="gpu",
choices=["gpu", "cpu", "xpu"],
choices=["gpu", "cpu", "xpu", "npu"],
help="Device selected for inference.", )
parser.add_argument('--cpu_threads', type=int, default=1)

@ -591,7 +591,8 @@ def get_predictor(
config = inference.Config(
str(Path(model_dir) / model_file), str(Path(model_dir) / params_file))
config.enable_memory_optim()
if device != "npu":
config.enable_memory_optim()
config.switch_ir_optim(True)
if device == "gpu":
config.enable_use_gpu(100, device_id)

Loading…
Cancel
Save