python torch使用gpu cuda

windows 安装cuda

注意选择合适的版本

https://developer.nvidia.com/cuda-toolkit-archive  

符合以下版本的可直接复制地址

windows pip cuda 11.8

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

安装pytorch-cuda

https://pytorch.org/

验证是否安装成功

新建.py文件，运行以下代码

import torch

if __name__ == '__main__':
    print("torch版本:", torch.__version__)
    print("CUDA是否可用", torch.cuda.is_available())
    print("可用的CUDA数量:", torch.cuda.device_count())
    print("CUDA的版本号", torch.version.cuda)

输出样例

torch版本: 2.0.0+cu118
CUDA是否可用 True
可用的CUDA数量: 1
CUDA的版本号 11.8