From d4c0517fdcb307b160b6070812f578bfb9e51923 Mon Sep 17 00:00:00 2001
From: ronny1996 <524019753@qq.com>
Date: Thu, 2 Dec 2021 12:48:00 +0000
Subject: [PATCH] fix bug for rocm

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

diff --git a/tools/infer/utility.py b/tools/infer/utility.py
index f437056e..af50c5e6 100644
--- a/tools/infer/utility.py
+++ b/tools/infer/utility.py
@@ -311,7 +311,10 @@ def create_predictor(args, mode, logger):
 
 
 def get_infer_gpuid():
-    cmd = "env | grep CUDA_VISIBLE_DEVICES"
+    if not paddle.fluid.core.is_compiled_with_rocm():
+        cmd = "env | grep CUDA_VISIBLE_DEVICES"
+    else:
+        cmd = "env | grep HIP_VISIBLE_DEVICES"
     env_cuda = os.popen(cmd).readlines()
     if len(env_cuda) == 0:
         return 0
-- 
GitLab