RuntimeError: CUDA error (10): invalid device ordinal

最新推荐文章于 2024-08-23 09:44:21 发布

This is bill

最新推荐文章于 2024-08-23 09:44:21 发布

阅读量1.2k

点赞数

分类专栏：机器学习 python

本文链接：https://blog.csdn.net/Scythe666/article/details/88529725

版权

机器学习同时被 2 个专栏收录

321 篇文章 17 订阅

订阅专栏

python

97 篇文章 6 订阅

订阅专栏

两张卡跑的train，在本机一张卡上test不了
结果很简单，直接将 model.module 拿出来就行了
因为之前双卡用的parallel

def __init__(self, module, device_ids=None, output_device=None, dim=0):
        super(DataParallel, self).__init__()

        if not torch.cuda.is_available():
            self.module = module
            self.device_ids = []
            return

        if device_ids is None:
            device_ids = list(range(torch.cuda.device_count()))
        if output_device is None:
            output_device = device_ids[0]