PyTorch 1.0.x - 支持 CUDA 7.5
PyTorch 1.1.x - 支持 CUDA 8.0
PyTorch 1.2.x - 支持 CUDA 9.0
PyTorch 1.3.x - 支持 CUDA 9.2
PyTorch 1.4.x - 支持 CUDA 10.1
PyTorch 1.5.x - 支持 CUDA 10.2
PyTorch 1.6.x - 支持 CUDA 11.0
PyTorch 1.7.x - 支持 CUDA 11.0/11.1
PyTorch 1.8.x - 支持 CUDA 11.1/11.2
PyTorch 1.9.x - 支持 CUDA 11.1/11.2/11.3
PyTorch 1.10.x - 支持 CUDA 11.1/11.2/11.3/11.4
 

https://pytorch.org/get-started/locally/

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118