Solved by RonMcKay in post #2 
              
              
                I solved it myself by modifying the code for the DataParallel wrapper. Now the predicted tensor is broadcasted to the different gpus instead of scattered. The sampled outputs are then accumulated on each gpu and gathered all at once when enough samples were reached. If you have memory constraints yo…