Solved by spanev in post #2 
              
              
                Hi @orena1, 
First, in autograd terms, leaf nodes are the input of the forward (the last  on which backpropagation will be applied) and the root the output of the forward network. 
Remember that when you use MSELoss, the last operation to be applied is a reduction (by default Mean but you can also c…