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分享的方法,感兴趣的链接中还有另一种方法。