Ma/video-subtitle-remover
1
0
Fork 0
mirror of https://github.com/YaoFANGUK/video-subtitle-remover.git synced 2026-02-26 21:54:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

优化 get_infer_gpuid 函数，增加对 PaddlePaddle 版本的检查

Browse Source
This commit is contained in:
xshaitt
2025-04-18 11:22:02 +08:00
parent 9f9cded1ff
commit 5ac9f67ad5
1 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
backend/tools/infer/utility.py
View File

@@ -340,7 +340,14 @@ def get_infer_gpuid():
if sysstr == "Windows":
return 0
if not paddle.fluid.core.is_compiled_with_rocm():
# 检查 PaddlePaddle 版本
paddle_version = paddle.__version__
major_version = int(paddle_version.split('.')[0])
if major_version >= 2:
is_rocm = paddle.device.is_compiled_with_rocm()
else:
is_rocm = paddle.fluid.core.is_compiled_with_rocm()
if not is_rocm:
cmd = "env | grep CUDA_VISIBLE_DEVICES"
else:
cmd = "env | grep HIP_VISIBLE_DEVICES"