diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 527a811d620efac33ece9cdbd4b6196e18a8497d..707328f28468db86c5061795d04713dc3b21a5cb 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -159,6 +159,11 @@ def create_predictor(args, mode, logger): precision = inference.PrecisionType.Float32 if args.use_gpu: + gpu_id = get_infer_gpuid() + if gpu_id is None: + raise ValueError( + "Not found GPU in current device. Please check your device or set args.use_gpu as False" + ) config.enable_use_gpu(args.gpu_mem, 0) if args.use_tensorrt: config.enable_tensorrt_engine(