未验证 提交 09604c38 编写于 作者: M MissPenguin 提交者: GitHub

Merge pull request #4824 from ronny1996/cp_pr4823

cherry-pick #4823
......@@ -20,6 +20,7 @@ import numpy as np
import json
from PIL import Image, ImageDraw, ImageFont
import math
import paddle
from paddle import inference
import time
from ppocr.utils.logging import get_logger
......@@ -286,11 +287,17 @@ def create_predictor(args, mode, logger):
def get_infer_gpuid():
cmd = "nvidia-smi"
if not paddle.fluid.core.is_compiled_with_rocm():
cmd = "nvidia-smi"
else:
cmd = "rocm-smi"
res = os.popen(cmd).readlines()
if len(res) == 0:
return None
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
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册