目录

2024-12-23-pytorch测试GPU是否可用

JAY.LIN 收录于 Python

2024-12-23 约 63 字预计阅读 1 分钟

https://bing.ee123.net/img/rand?artid=89709116

目录

pytorch：测试GPU是否可用

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

True
cuda:0
GeForce GTX 1080
tensor([[0.9530, 0.4746, 0.9819],
        [0.7192, 0.9427, 0.6768],
        [0.8594, 0.9490, 0.6551]], device='cuda:0')

68747470733a2f2f626c6f67:2e6373646e2e6e65742f77656978696e5f3335353736383831:2f61727469636c652f64657461696c732f3839373039313136