查看torch版本是否为GPU版本
import torch
print(torch.__version__) //我的是CPU版本
如果不是GPU版本下载GPU版本
-
查看cuda版本
nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:41:10_Pacific_Daylight_Time_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0 //我的是11.8版本 -
下载与cuda版本一致的pytorch
//下载pytorch 11.8版本，去pytorch官网生成11.8对应的pytorch版本
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118