From 5bb6ae05c4f2c3c22a17ecf4b227164932cd6a59 Mon Sep 17 00:00:00 2001
From: LDOUBLEV <liuvv0203@outlook.com>
Date: Tue, 17 Aug 2021 02:45:26 +0000
Subject: [PATCH] add rec get_gpu_id

---
 tools/infer/predict_rec.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tools/infer/predict_rec.py b/tools/infer/predict_rec.py
index bb4a3170..97dfa521 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'
                 ],
-- 
GitLab