fix bug for rocm

2021-12-02 12:48:00 +00:00 · 2021-12-02 12:48:00 +00:00 · d4c0517fdc
parent c81baca7d9
commit d4c0517fdc
1 changed files with 4 additions and 1 deletions
--- 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