使用!nvidia-smi在jupyter notebook中查看cuda版本。
cuda版本为11.2,上官网https://download.pytorch.org/whl/torch_stable.html上寻找对应版本torch,发现没有11.2,可以考虑下载cuda11.1版本相对应的torch。
也可以在https://pytorch.org/get-started/previous-versions/上直接使用代码下载
# CUDA 11.1
pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html
最终得到如下结果