How does one use Pytorch (+ cuda) with an A100 GPU?

From the link pytorch site from @SimonB ‘s answer, I did:

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

This solved the problem for me.