1. 查看驱动程序版本
方法一: 桌面右击 选择NVIDIA控制面板 选择帮助中的系统信息,查看驱动程序版本。
方法二: 终端输入
nvidia-smi
2. 查询匹配的cuda版本
去https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html查询匹配的CUDA版本,我的是10.2.89
3. 安装pytorch
pytorch官网https://pytorch.org/
选择如图,使用pip安装
pip3 install torch==1.8.1+cu102 torchvision==0.9.1+cu102 torchaudio===0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
4. 验证安装效果
import torch
torch.cuda.is_available()