代码
import torch
flag = torch.cuda.is_available()
print(flag)
ngpu= 1
# Decide which device we want to run on
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda())
运行结果如下表示GPU可用了
True
cuda:0
NVIDIA GeForce RTX 3060
tensor([[0.7567, 0.0458, 0.5618],
[0.6300, 0.2157, 0.8634],
[0.2069, 0.5930, 0.2276]], device='cuda:0')
