确保环境可以运行cuda版本的pytorch

1.使用nividia-smi确保有cuda。

根据cuda版本号选择对应的pytorch。

Previous PyTorch Versions | PyTorch

pip3 install torch==1.10.1+cu102 torchvision==0.11.2+cu102 -f https://download.pytorch.org/whl/cu102/torch_stable.html

你可以通过以下Python代码测试CUDA是否可以正常工作：

python
import torch
torch.cuda.is_available()

返回True，则说明CUDA可以正常工作。