Hi,
You can use tensor.register_hook(your_fn) to be able to print the gradient for each given Tensor.
In particular, here, I think that the loss before the nansum won’t have any nan as you expect.
The problem is that the backward of the mse will create nan values because some of the y has nan. You will need to filter these out before computing the MSE to avoid nan in the gradients flowing back.