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 it would be a neat way to implement what you want as well: make the .grad  fields the same so the two backward passes actually accumulate in the same Tensor. And divide your learning rate by 2 so that you still do the same update (assuming simple SGD).