Ao instalar a tocha, preste atenção na versão cuda da placa gráfica.
Por exemplo, no mesmo ambiente do RTX2080, o programa pode rodar normalmente, mas ao alternar para o A100, o seguinte erro será reportado:
NVIDIA A100-PCIE-40GB com capacidade CUDA sm_80 não é compatível com a instalação atual do PyTorch. A instalação atual do PyTorch suporta recursos CUDA sm_37 sm_50 sm_60 sm_70 sm_75. Se você deseja usar a GPU NVIDIA A100-PCIE-40GB com PyTorch, verifique as instruções em Iniciar localmente | PyTorch
Provavelmente significa: NVIDIA A100-PCIE-40GB tem um poder de computação CUDA de 8.0, que não corresponde à versão existente do PyTorch. O poder de computação CUDA suportado pela versão existente do PyTorch é 3.7, 5.0, 6.0, 7.0, 7.5.
O poder de computação do CUDA suportado está relacionado à versão do cuda instalada. O CUDA 10.2 suporta apenas o poder de computação 3.7, 5.0, 6.0, 7.0 e não suporta o poder de computação 8.0. E cuda11 suporta poder de computação 8.0.
A versão do torch atualmente instalada é 1.7.0, então você precisa instalar o cuda11 e superior, e uma versão que não entre em conflito com o torch 1.7.0.
Entre no site oficial do PyTorch : Versões anteriores do PyTorch | PyTorch
Selecione a versão CUDA apropriada, você também pode ir para Versões anteriores do PyTorch para ver a seleção,
Finalmente escolhi a versão da v1.7.1 CUDA 11.0
# CUDA 11.0
pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
problema resolvido.
Referência: https://zhuanlan.zhihu.com/p/427395039