It’s hard to tell what’s going on without a self-contained example.
Does the loss in the end require gradients?
If not, I would recommend to go through the steps in the forward and print if it has grad.
If it does, you can add backward hooks that print:
def add_detect_backward(name, tensor):
  def bwh(x):
     print("Computing backward of ", name)
  tensor.register_hook(bwh)

a = torch.randn(5, 5, requires_grad=True)
add_detect_backward("a", a)
b = 5 * a
add_detect_backward("b", b)
c = b.sum()
c.backward()

This should give you a good picture of where the gradient reaches and where you drop it accidentally.