I’ll try to figure out how to reproduce it with a smaller example, right now the code is pretty big.
However the only differences that I have between multi and single gpu are:
if args.gpu == 99:
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
else:
    device = torch.device("cuda:{}".format(args.gpu) if torch.cuda.is_available() else "cpu")


and this
if args.gpu == 99:
    net = nn.DataParallel(net)
net.to(device)

basically if I set the arg to 99 it’ll use dataparallel. Not sure what I’m missing.