@Xiaofeng_Wu As far as I understand you are using B(A(x)) as a “fixed” target, i.e. B doesn’t need gradients and shouldn’t be trained. In my opinion, detaching should be alright, but I’m not sure if there are some pitfalls I’m missing using KLDiv.
If you don’t detach A(x) while passing it to B, could you check the gradients in A?
@viraat I think the code looks good. I just re-created a small dummy experiment and I get valid gradients for encoder and classifier, while adversary doesn’t have any gradients.
Could you please double check it by calling backward on both losses separately and checking the gradients of your models?