模型训练前检查是否可以成功使用GPU训练，附上国内下载pytorch方法

每天八杯水D

于 2024-12-04 09:15:00 发布

阅读量353

点赞数 8

分类专栏： pytorch 文章标签：深度学习 python pytorch cuda

版权

2 篇文章

订阅专栏

代码

import torch

# 检查是否有可用的 GPU
if torch.cuda.is_available():
    print("Training on GPU.")
else:
    print("Training on CPU.")

pytorch官网：PyTorch

如上图红色下划线部分所示为外网源，只需删掉下面红色部分链接，切换为清华源即可。

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

pip3 install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple