pytorch 使用GPU进行神经网络训练出现问题 RuntimeError: CUDA error: invalid device ordinal

最新推荐文章于 2024-06-05 15:23:31 发布

Childhood_Sweetheart

最新推荐文章于 2024-06-05 15:23:31 发布

阅读量2.2k

点赞数

文章标签：深度学习

本文链接：https://blog.csdn.net/Childhood_Sweetheart/article/details/115165134

版权

实验室里有两块GTX1080Ti 的显卡，一开始只是使用第一块显卡，需要在代码中加入这样一行代码：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

但是第一块显卡被使用了，要用第二块显卡，于是我就将代码改为：

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

但是出现了这样的错误，在网上查阅了资料才发现需要这样改：

import os 
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu' )

需要注意的是 os.environ['CUDA_VISIBLE_DEVICES'] = '1' 这行代码一定要放在 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu' ) 的前面。

关注