首先,查看显卡驱动版本:
nvidia-smi
出不来结果的,请先安装显卡驱动(https://developer.nvidia.com/cuda-downloads)。
安装torch:
地址:https://pytorch.org/get-started/locally/
可使用在线安装,按需选择,然后生成脚本,比如:
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
或者下载whl后安装(https://download.pytorch.org/whl/torch/)。
按自己的显卡需要,尽可能安装最新的(有时间显卡并没有最新的驱动,只需要装一个能支持的、最新的就好)。
注意,要选择带 +cu的,就是这样的:torch-2.1.1+cu121-cp310-cp310-win_amd64.whl