【pytorch】torch1.8.1安装、查看torch版本、GPU是否可用

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