首先说明,很多博主文章里没有提到的一点,pytorch包和torchvision包是一同下载的,且pytroch包包含在torchvision内,因此要先把环境中的pytorch,torchvision,torchaudio都一一卸载掉
使用命令
pip uninstall torchvision pytorch torchaudio torch
之后用nvidia-smi查看对应cuda 版本
然后使用这个命令
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
后面的113换成你的cuda版本对应的(我的是11.3所以就是113)
这样下载下来的就不是cpu版本的了
我们来check一下,命令行键入
python
import torch
print(torch.cuda.is_available())