Solved by ptrblck in post #13 
              
              
                Looks good to me. To search for the problematic part, could you repeat this procedure with random tensors as input, i.e. don’t use your Dataset and DataLoader? 
Since you are seeding, the random tensor should be the same in each run.