pytorch判断是否安装了正确的cuda环境

#判断是否安装了cuda
import torch

print(f"cuda可用性：{torch.cuda.is_available()}")

print("GPU names:")
for i in range(torch.cuda.device_count()):
    print(f"{i} : {torch.cuda.get_device_name(i)}")

# print(torch.cuda.current_device())


# 判断是否安装了cuDNN
from torch.backends import cudnn

print(f'是否安装cuDNN：{cudnn.is_available()}')  # 返回True则说明已经安装了cuDNN


# 看是否支持sm_86
torch.zeros(1).cuda()

# 进一步验证
dev = "cuda:0"
device = torch.device(dev)
a = torch.zeros(4,3)
a = a.to(device)
a