Besides that, I also tried to use hook to monitor the gradients.
Since method 2 suggest the issue is with:
logits = F.linear(F.normalize(emb), F.normalize(weight))
I defined a tensor hook to check is gradients of emb and weight contain Nan. The strange thing is that, when I turn on the hook, anomaly detection gives no error. But when I turn off the hook, anomaly detection throws runtime error at the same iterations.
Is there any conflict between hooks and anomaly detection?