如何将Pytork（+cuda）与A100 GPU配合使用？

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上使用

Answer 2

我也有同样的问题。您需要安装CUDA 11.0而不是10.2，并为此CUDA版本重新安装PyTorch

Answer 3

从@SimonB的答案中的链接pytorch site，我做到了：

pip3 install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html

这为我解决了问题