torch 测试GPU能否正常使用

运行程序

import torch
print(torch.cuda.is_available())

num_gpu= 1
# Decide which device we want to run on
device = torch.device("cuda:0" if (torch.cuda.is_available() and num_gpu > 0) else "cpu")
print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda())

输出结果：

True
cuda:0
GeForce RTX 3090
tensor([[0.7202, 0.7553, 0.8924],
        [0.6794, 0.1734, 0.9777],
        [0.3063, 0.4634, 0.0939]], device='cuda:0')

Process finished with exit code 0