Solved by yuval in post #5 
              
              
                I managed to solve the problem. 
It is related to the fact that my code used: 
           for i in range(len(train_loader)):
               X, y1, y2 = train_loader.__iter__().__next__()
               <process batch>

Replacing it with the following solved the problem: 
           for i, (X, y1, y2…