diff --git a/tools/infer/predict_rec.py b/tools/infer/predict_rec.py index bb4a31706471b9b1745519ac9f390d01b60d5d44..97dfa5214628123d0c9b7edd7d94060a2bfd2a1e 100755 --- a/tools/infer/predict_rec.py +++ b/tools/infer/predict_rec.py @@ -68,6 +68,7 @@ class TextRecognizer(object): if args.benchmark: import auto_log pid = os.getpid() + gpu_id = utility.get_infer_gpuid() self.autolog = auto_log.AutoLogger( model_name="rec", model_precision=args.precision, @@ -77,7 +78,7 @@ class TextRecognizer(object): inference_config=self.config, pids=pid, process_name=None, - gpu_ids=0 if args.use_gpu else None, + gpu_ids=gpu_id if args.use_gpu else None, time_keys=[ 'preprocess_time', 'inference_time', 'postprocess_time' ],