You could  exploit the fact that nan != nan to identify the tensor that first contains a nan. To illustrate what I mean by nan != nan:
>>> import torch
>>> import numpy as np
>>> t = torch.tensor([1., 2., np.nan])
>>> t != t
 tensor([ 0,  0,  1], dtype=torch.uint8)

Then you can check in your code which variable via if (t != t).any():
for each tensor you have in the graph. It’s probably best to use a debugger for that. E.g., raise an error when if (t != t).any(): evaluates to True and then expect all the current variables.