python - 如何在 A100 GPU 上使用 Pytorch (+ cuda)？

Question

我试图将我当前的代码与 A100 gpu 一起使用，但出现此错误:

---> backend='nccl'
/home/miranda9/miniconda3/envs/metalearningpy1.7.1c10.2/lib/python3.8/site-packages/torch/cuda/__init__.py:104: UserWarning: 
A100-SXM4-40GB with CUDA capability sm_80 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_37 sm_50 sm_60 sm_61 sm_70 sm_75 compute_37.
If you want to use the A100-SXM4-40GB GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/

这确实令人困惑，因为它指向通常的 pytorch 安装，但没有告诉我要用于我的特定硬件 (A100) 的 pytorch 版本 + cuda 版本的哪个组合。为 A100 安装 pytorch 的正确方法是什么？

这些是我尝试过的一些版本:

# conda install -y pytorch==1.8.0 torchvision cudatoolkit=10.2 -c pytorch
# conda install -y pytorch torchvision cudatoolkit=10.2 -c pytorch
#conda install -y pytorch==1.7.1 torchvision torchaudio cudatoolkit=10.2 -c pytorch -c conda-forge
# conda install -y pytorch==1.6.0 torchvision cudatoolkit=10.2 -c pytorch
#conda install -y pytorch==1.7.1 torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge

# conda install pytorch torchvision torchaudio cudatoolkit=11.0 -c pytorch
# conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge
# conda install -y pytorch torchvision cudatoolkit=9.2 -c pytorch # For Nano, CC
# conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge

请注意，这可能很微妙，因为我过去在这台机器 + pytorch 版本上遇到过这个错误:
How to solve the famous `unhandled cuda error, NCCL version 2.7.8` error?

Answer 1

我有一台 A100 并且已经成功

conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c nvidia

pytorch site 上现在也推荐了