Solved by ptrblck in post #6 
              
              
                Thanks for the update. 
The workaround is invalid, as you would call the underlying model on the default device and would skip the nn.DataParallel wrapper. 
The device mismatch is raised, since torch.distributions do not have a to() method and are not registered as modules. 
Here is a minimal code s…