Unverified Commit 0ba1f76d authored by Yao Matrix's avatar Yao Matrix Committed by GitHub
Browse files

enable print_env on xpu (#11507)



* detect xpu in print_env
Signed-off-by: default avatarYAO Matrix <matrix.yao@intel.com>

* enhance code, test passed on XPU
Signed-off-by: default avatarYao Matrix <matrix.yao@intel.com>

---------
Signed-off-by: default avatarYAO Matrix <matrix.yao@intel.com>
Signed-off-by: default avatarYao Matrix <matrix.yao@intel.com>
parent d6bf268a
......@@ -34,13 +34,24 @@ try:
print("Torch version:", torch.__version__)
print("Cuda available:", torch.cuda.is_available())
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())
if torch.cuda.is_available():
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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment