Merge pull request #4823 from ronny1996/fix_rocm

Fix bug for rocm

Merge pull request #4823 from ronny1996/fix_rocm
Fix bug for rocm
e16260c9 · MissPenguin · GitHub · 44ba6030 · d4c0517f · e16260c9
Unverified Commit e16260c9 authored Dec 03, 2021 by MissPenguin Committed by GitHub Dec 03, 2021
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

tools/infer/utility.py tools/infer/utility.py +4 -1

No files found.
--- a/tools/infer/utility.py
+++ b/tools/infer/utility.py
@@ -311,7 +311,10 @@ def create_predictor(args, mode, logger):


 def get_infer_gpuid():
+    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