当前位置:主页 > 软件编程 > Python代码 >

解决pytorch报错:AssertionError: Invalid device id的问题

时间:2021-01-12 13:20:49 | 栏目:Python代码 | 点击:

在服务器上训练的网络放到本地台式机进行infer,结果出现报错:

AssertionError: Invalid device id

仔细检查后发现原来服务器有多个GPU,当时开启了两个进行加速运算。

net1 = nn.DataParallel(net1, device_ids=[0, 1])

而本地台式机只有一个GPU,调用数量超出所以报错。

改为

net1 = nn.DataParallel(net1, device_ids=[0])

问题解决。

您可能感兴趣的文章:

相关文章