From 0ba1f76d4dde6d25b33dbdca73b6aa21bb682c56 Mon Sep 17 00:00:00 2001 From: Yao Matrix Date: Fri, 9 May 2025 19:00:51 +0800 Subject: [PATCH] enable print_env on xpu (#11507) * detect xpu in print_env Signed-off-by: YAO Matrix * enhance code, test passed on XPU Signed-off-by: Yao Matrix --------- Signed-off-by: YAO Matrix Signed-off-by: Yao Matrix --- utils/print_env.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/utils/print_env.py b/utils/print_env.py index 0a1cfbef13..2d2acb59d5 100644 --- a/utils/print_env.py +++ b/utils/print_env.py @@ -34,13 +34,24 @@ try: print("Torch version:", torch.__version__) print("Cuda available:", torch.cuda.is_available()) - print("Cuda version:", torch.version.cuda) - print("CuDNN version:", torch.backends.cudnn.version()) - print("Number of GPUs available:", torch.cuda.device_count()) if torch.cuda.is_available(): + print("Cuda version:", torch.version.cuda) + print("CuDNN version:", torch.backends.cudnn.version()) + print("Number of GPUs available:", torch.cuda.device_count()) device_properties = torch.cuda.get_device_properties(0) total_memory = device_properties.total_memory / (1024**3) print(f"CUDA memory: {total_memory} GB") + + print("XPU available:", hasattr(torch, "xpu") and torch.xpu.is_available()) + if hasattr(torch, "xpu") and torch.xpu.is_available(): + print("XPU model:", torch.xpu.get_device_properties(0).name) + print("XPU compiler version:", torch.version.xpu) + print("Number of XPUs available:", torch.xpu.device_count()) + device_properties = torch.xpu.get_device_properties(0) + total_memory = device_properties.total_memory / (1024**3) + print(f"XPU memory: {total_memory} GB") + + except ImportError: print("Torch version:", None)