Thanks for the reply! In Approach 2, the loss depends only on output[1,:], the elements of which are non-NaN, therefore the loss is non-NaN as well. Wouldn’t we expect the gradients to propagate properly in this case since the NaN tensor elements are “irrelevant” - they don’t feed into the loss function?
Or is it the case that if at least one component of a tensor is NaN, then the gradients can’t flow through this tensor?