pytorch多块gpu使用总结，报错AssertionError: Invalid device id

 

pytorch关于多块gpu使用总结，报错AssertionError: Invalid device id

实验室有4块gpu

想使用标号2，3块时报错：AssertionError: Invalid device id

解决方案：

修改pytorch感受的设备编号。

os.environ['CUDA_VISIBLE_DEVICES'] = '2,3'

cuda编号还是从0开始。 

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

 修改pytorch设备编号还为0，1。

 if torch.cuda.device_count() > 1:
        print(torch.cuda.device_count())
        net = nn.DataParallel(net, device_ids=[0, 1])
    net.to(device)

感谢https://blog.csdn.net/kongkongqixi/article/details/100521590分享的方法，感兴趣的链接中还有另一种方法。