Solved by ptrblck in post #16 
              
              
                It seems the error is raised by rewrapping the model into nn.DataParallel in each iteration. 
Move contact_net = torch.nn.DataParallel(contact_net) before the epoch loop and it should work. 
I don’t know, why this usage gives a device mismatch error and think a better error message should be raised.…