Just updated to the latest torch 1.8.1 and now the error has changed to:
....python3.7/site-packages/torch/cuda/__init__.py:104: UserWarning:
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.
If you want to use the A100-PCIE-40GB GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/

I had a look at the link provided, but I am not very sure how I can update or change my pytorch installation to allow sm_80.