OK I see. Thanks for the detailed explanations.
By saying “split all of these into different lines”, do you mean this?
norm_emb = F.normalize(emb)
norm_w = F.normalize(weight)
logits = F.linear(norm_emb, norm_w)

I checked the values for original and normalized emb/weights, and didn’t find any problem. But according to hook, the gradients of weight are nan (not all nan, only some rows), after that the training still goes on successfully for several iterations, and then gradients of weight happen repeatedly, leading to Nan loss.
The feedback given by hook and anomaly detection are different, gradient of weights is nan VS normalization of weights has issue. I don’t know if they refer to the same problem or which one should I believe.