You should write your own backward if your function is not differentiable or you don’t want to use the autograd function. this will be the backward which will be called when you call loss.backward(). If you want individual grad’s like in your second case, maybe you call loss.backward() and then read the .grad attributes of A and B.
Just make sure .requires_grad attribute of all the tensors for which you want gradients is set to True