Hey,
If you don’t want gradients corresponding to headA, then yes I think the best thing to do is to just not backrprop for it.
Note that if you share the weights between A and B, then the weights in A will still change due to the update on B.