Linux x86_64 GA102 [GeForce RTX 3090]**安装pytorch-gpU 1.8.1 ,显卡驱动,cuda,cudnn

Linux x86_64 GA102 [GeForce RTX 3090]**安装pytorch-gpU 1.8.1 ,显卡驱动,cuda,cudnn

 pip install torchaudio-0.8.1

 

pip install torchvision==0.9.1

验证

import torch
import torch.nn as nn
import torchvision


torch.cuda.is_available()  ## 输出应该是True

print(torch.__version__)
print(torch.version.cuda)
print(torch.backends.cudnn.version())
print(torch.cuda.get_device_name(0))

 
t=torch.nn.Linear(3,3)
t.to("cuda:0")
 
input=torch.randn((3,3)).requires_grad_().to("cuda:0")
output=t(input)
 
loss=torch.sum(output)
torch.autograd.grad(loss,input,retain_graph=True)  ## 输出应该是一个gpu上的梯度矩阵
loss.backward()