在conda虚拟环境下安装torch==1.8.1+gpu版本。
本机环境
- CUDA 11.1
- python 3.7
安装torch1.8.1
使用在线安装即可,离线安装需要下载对应的torch、torchvision、torchaudia安装包有点麻烦。
1. 在官网找到对应cuda版本
2. 安装命令
pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
如果出现找不到版本,请使用pip3 install
查看版本
import torch
# 检测torch、cuda、cudnn版本
print(torch.__version__)
print(torch.version.cuda)
print(torch.backends.cudnn.version())
#是否可用gpu
flag = torch.cuda.is_available()
print(flag)