驗證pytorch是否安裝成功

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