From 4a9e50e1eb1b41cb01d2eb766aa82d9123b03b6b Mon Sep 17 00:00:00 2001
From: Jim <tianwai263@gmail.com>
Date: Mon, 27 Dec 2021 12:18:23 +0800
Subject: [PATCH] fix  #5065

---
 tools/infer/utility.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tools/infer/utility.py b/tools/infer/utility.py
index 21bbee09..33ed6212 100644
--- a/tools/infer/utility.py
+++ b/tools/infer/utility.py
@@ -313,6 +313,11 @@ def create_predictor(args, mode, logger):
 
 
 def get_infer_gpuid():
+    if os.name == 'nt':
+        try:
+            return int(os.environ['CUDA_VISIBLE_DEVICES'].split(',')[0])
+        except KeyError:
+            return 0
     if not paddle.fluid.core.is_compiled_with_rocm():
         cmd = "env | grep CUDA_VISIBLE_DEVICES"
     else:
-- 
GitLab