本地下载并安装pytorch

查看cuda版本

nvcc -V
nvidia-smi

下载所需的torch torchvision torchaudio

地址：https://download.pytorch.org/whl/torch_stable.html

安装

先进入虚拟环境，再进入对应的安装目录

pip install (XXX.whl)
# 先安装torch，再安装torchvision，再安装torchaudio

验证

(sc_yolov5) [zqchen@gpurtx02 environment]$ python
Python 3.9.0 (default, Nov 15 2020, 14:28:56) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available
<function is_available at 0x7f7c9cfb8670>
>>> print(torch.__version__)
1.12.1+cu113
>>> print(torch.cuda.is_available())
True