diff --git a/tools/infer/utility.py b/tools/infer/utility.py index f437056ec7b10e28e626d2028b6401cebc647bb1..af50c5e6a8cb39faf416dbe7adb516c4db05aef5 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