riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-03 04:40:06 +06:00
Code Issues Actions 28 Packages Projects Releases Wiki Activity

add XPU info print in print_env (#38282)

Browse Source 
Signed-off-by: Matrix Yao <matrix.yao@intel.com>
This commit is contained in:
Yao Matrix 2025-05-22 17:03:56 +08:00 committed by GitHub
parent b369a65480
commit b26157d64c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 16 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
utils/print_env.py
View File

@ -21,6 +21,7 @@ import os
import sys
import transformers
from transformers import is_torch_xpu_available
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@ -32,11 +33,21 @@ try:
import torch
print("Torch version:", torch.__version__)
print("Cuda available:", torch.cuda.is_available())
accelerator = "NA"
if torch.cuda.is_available():
accelerator = "CUDA"
elif is_torch_xpu_available():
accelerator = "XPU"
print("Torch accelerator:", accelerator)
if accelerator == "CUDA":
print("Cuda version:", torch.version.cuda)
print("CuDNN version:", torch.backends.cudnn.version())
print("Number of GPUs available:", torch.cuda.device_count())
print("NCCL version:", torch.cuda.nccl.version())
elif accelerator == "XPU":
print("SYCL version:", torch.version.xpu)
print("Number of XPUs available:", torch.xpu.device_count())
except ImportError:
print("Torch version:", None)