一、检测脚本
import torchprint(torch.__version__)
print(torch.cuda.is_available())
print(torch.version.cuda)
输出
二、解决排查
1、安装CUDA
根据驱动选择要安装的版本
安装网站:CUDA Toolkit Archive | NVIDIA Developer
安装成功后
2、安装cuDNN
注意:需要注册才能进行安装
具体网站:cuDNN 9.1.1 Downloads | NVIDIA Developer
下载后,需要解压
需要将bin、include、和lib/x86目录中的文件拷贝到亲们cuda安装的对应目录中
并且新增如下环境变量:
配置成功后,使用nvidia-smi命令检查如下
3、修改pytorch版本,从cpu版本改为gpu版本
1)卸载cpu版本
pip3 uninstall torch torchvision torchaudio
2)安装gpu版本
网站:Start Locally | PyTorch
使用网站上面生成的命令进行安装
三、验证成果
执行文章最开始的脚本,结果如下