import torch
#查看torch版本
print(torch.version)
#查看cuda版本
print(torch.version.cuda)
#GPU是否可用
print(torch.cuda.is_available())
#返回gpu数量
print(torch.cuda.device_count())
#返回gpu名字,设备索引默认从0开始
print(torch.cuda.get_device_name(0))
#例子
print(torch.rand(3,3).cuda())
运行结果实例
1.5.0
10.2
True
1
NVIDIA GeForce RTX 3060 Laptop GPU
tensor([[0.7567, 0.0458, 0.5618],
[0.6300, 0.2157, 0.8634],
[0.2069, 0.5930, 0.2276]], device='cuda:0')