Hi,
I think the simplest way to understand what will happen here is to know that the autograd lives below torch.nn and is completely unaware of what torch.nn does.
So in this case, whatever is the Tensor you give to the rest of the net is the one that will get gradients (it does not matter if it comes from a hook or not).
And in this case, since A_hooked depends on A, then the gradients will flow back from A_hooked to A.