import torch

# 检查 CUDA 是否可用
print("CUDA available:", torch.cuda.is_available())

# 如果 CUDA 可用，打印当前设备信息
if torch.cuda.is_available():
    print("Current CUDA device:", torch.cuda.current_device())
    print("Name of the CUDA device:", torch.cuda.get_device_name(0))
    print("Number of CUDA devices:", torch.cuda.device_count())
else:
    print("CUDA is not available. PyTorch will use CPU instead.")