Solved by albanD in post #2 
              
              
                Hi, 
I would bet the issue is with: 
        pA.grad = avg
        pB.grad = avg

Here you set the same Tensor to be the gradient for both parameters. So the backward passes will accumulate in the Same Tensor. 
If you don’t want that, you need to add a .clone() for at least one of them. 
Note that i…