pytorch 验证GPU是否可用

2022-10-09 10:26:25

代码

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')

  • 作者:小猫不会去楼兰捉虫
  • 原文链接:https://blog.csdn.net/anglemanyi/article/details/124766488
    更新时间:2022-10-09 10:26:25