When installing torch, be sure to pay attention to the cuda version of the graphics card.
For example, in the same environment on the RTX2080, the program can run normally, but when you switch to the A100, the following error will be reported:
NVIDIA A100-PCIE-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_70 sm_75. If you want to use the NVIDIA A100-PCIE-40GB GPU with PyTorch, please check the instructions at Start Locally | PyTorch
The approximate meaning is: NVIDIA A100-PCIE-40GB has a CUDA computing power of 8.0, which does not match the existing PyTorch version. The CUDA computing power supported by the existing PyTorch version is 3.7, 5.0, 6.0, 7.0, 7.5.
The supported CUDA computing power is related to the version of cuda installed. CUDA 10.2 only supports 3.7, 5.0, 6.0, 7.0 computing power, and does not support 8.0 computing power. And cuda11 supports 8.0 computing power.
The currently installed torch version is 1.7.0, so you need to install cuda11 and above, and a version that does not conflict with torch 1.7.0.
Enter the PyTorch official website : Previous PyTorch Versions | PyTorch
Select the appropriate CUDA version, you can also go to Previous PyTorch Versions to view the selection,
Finally chose the version of 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
problem solved.
Reference: https://zhuanlan.zhihu.com/p/427395039