After Further debugging, I find that add a gradient hook to vs and modify the gradient to replace the nan with 0 does  solve the problem mentioned above. That is to say, the nan gradient from torch.std() is replaced with 0.
However, I then found there is another nan bug in this code. And since I’m using torch.autograd.detect_anomaly() to find out which line is the culprit, the line
vs_std = torch.std(vs, dim=(2, 3)) + self.eps
always shows up in the error log. It seems that torch.autograd.detect_anomaly() still gives an error at this line even after adding a gradient hook to vs (in fact, I have to comment out torch.autograd.detect_anomaly() to make the code run normally). And that covers up the trace of the second nan bug.
What a code snippet! I never thought of another nan bug, so naturally I’ve wasted lots of time.
Here is the thing: I’ve added gradient hooks to every module in this custom block, and I find the linear layers are also having nan gradient. There are many linear layers receiving nan grad_out and sending grad_in, but there is also this one ‘linear’ layer ‘fc_w’ receving normal ‘grad_out’ and sending nan ‘grad_in’.  At first, I thought it was the input es_domain was nan. But then I found out es_domain became nan at iteration 1, while nan gradient in linear layer happened at iteration 0.
So I wonder how to fix this bug? Or is there any way to inpect the internal backprop procedure in a linear layer?
Thanks again!