Thanks! Thomas
That is indeed the case for me. I checked the vs_std tensor and found something like vs[14, 205, :, :] was an all-zero tensor. So I assmue modify the nan gradient of vs with 0 would solve the problem.
However, I’ve tried to add gradient hook to both vs and vt (or vs_std and vt_std), and it didn’t work.
def hook_fn_backward(grad):
    grad = torch.where(torch.logical_not(torch.isnan(grad)), grad, torch.zeros_like(grad).to(grad.device))
return grad

vs.register_hook(hook_fn_backward)
vs_mean.register_hook(hook_fn_backward)
vs_std.register_hook(hook_fn_backward)

It still gives following error
RuntimeError: Function 'StdBackward1' returned nan values in its 0th output.
The strange thing is, I tried to set trace in hook_fn_backward, but I only found that torch.any(torch.isnan(grad)) == False.  So where should I add this hook exactly?