I’m not sure it’s working in the way you want it to.
Your architecture looks like a GAN, where only the generator will be trained. If you detach your prediction (with .data) you will lose the gradients for A.
Could you check if you get valid gradients in A?