Merge pull request #3728 from LDOUBLEV/fix_tests

raise error when run on cpu but set use_gpu as True
This commit is contained in:
dyning 2021-08-18 17:28:27 +08:00 committed by GitHub
commit ab16f2e4f9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 0 deletions

View File

@ -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(