I think you should also clone the gradients of A&B. Can you change like below. Without clone data is not explicitly copied to destination I guess.
for net1,net2 in zip(A.named_parameters(),B.named_parameters()):
                net2[1].data.grad = net1[1].data.grad.clone