riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-31 02:02:21 +06:00
Code Issues Actions 33 Packages Projects Releases Wiki Activity

Fix ROCm get_device_capability

Browse Source 
Signed-off-by: cyy <cyyever@outlook.com>
This commit is contained in:
cyy 2025-04-10 01:36:14 +08:00
parent ef38892ed4
commit 0e756434bd
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/transformers/testing_utils.py
View File

@ -3085,10 +3085,11 @@ def get_device_properties() -> DeviceProperties:
if IS_CUDA_SYSTEM or IS_ROCM_SYSTEM:
import torch
major, _ = torch.cuda.get_device_capability()
if IS_ROCM_SYSTEM:
major, _ = torch.hip.get_device_capability()
return ("rocm", major)
else:
major, _ = torch.cuda.get_device_capability()
return ("cuda", major)
elif IS_XPU_SYSTEM:
import torch