import torch
# 检查CUDA是否可用
if torch.cuda.is_available():
print("CUDA is available!")
# 还可以获取CUDA设备的数量
device_count = torch.cuda.device_count()
print(f"Number of CUDA devices: {device_count}")
# 获取第一块GPU的信息
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f"Device name: {torch.cuda.get_device_name(device)}")
# 或者进一步获取GPU的详细能力信息
capability = torch.cuda.get_device_capability(device)
print(f"Device capability: {capability}")
else:
print("CUDA is not available.")
返回
CUDA is available!
Number of CUDA devices: 4
Device name: NVIDIA GeForce RTX 2080 Ti
Device capability: (7, 5)