1 测试代码

import time
import torch

if __name__ == '__main__':
    print('torch版本：'+torch.__version__)
    print('cuda是否可用：'+str(torch.cuda.is_available()))
    print('cuda版本：'+str(torch.version.cuda))
    print('cuda数量:'+str(torch.cuda.device_count()))
    print('GPU名称：'+str(torch.cuda.get_device_name()))
    print('当前设备索引：'+str(torch.cuda.current_device()))

    device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")
    print(device)
    print(torch.rand(3, 3).cuda())

    for i in range(1,100000):
        start = time.time()
        a = torch.FloatTensor(i*100,1000,1000)
        a = a.cuda() #a = a
        a = torch.matmul(a,a)
        end = time.time() - start
        print(end)

结果：

torch版本：1.7.1
cuda是否可用：True
cuda版本：11.0
cuda数量:1
GPU名称：NVIDIA GeForce RTX 3090
当前设备索引：0
cuda:0
tensor([[0.3954, 0.7494, 0.8901],
        [0.6576, 0.1618, 0.6446],
        [0.1332, 0.2121, 0.8951]], device='cuda:0')

如果显示cuda可用，并且pytorch tensor可正常放到cuda上面去计算，那么就表明整个pytorch+cuda环境搭建成功，后面的for循环主要是为了测试GPU的显存是否能正常被pytorch调用。