Yes, it works!
After inserting some noise to both vs and new_vt as you suggested (without gradient hook), this code can now runs normally without giving nan values.
What I’m trying to do is to modify the moment from features of target domain (vt) with the moment from features of source domain (vs), and calculate some loss based on modified feature new_vt and its moment  new_et_domain. In practice, there are multiple domains and each domain has its corresponding custom block. That is to say, custom_block{i}.exchange(v_{i}, v_{j}) is performed on each pair of domains (for example, 3 domains in total).
However, there is one strange thing I don’t understand. If inserting noise works, then the nan should come from the std part. But why doesn’t gradient hook work as well? I tried to only add noise to vs, and it gives this error instead:
/opt/conda/conda-bld/pytorch_1573049306803/work/torch/csrc/autograd/python_anomaly_mode.cpp:57: UserWarning: Traceback of forward call that caused the error:
 new_vt = weight * vt + bias

RuntimeError: Function 'MulBackward0' returned nan values in its 0th output.

Also, if I add noise to vs and gradient hook to new_vt, the bias of fc_b becomes normal
 for name, param in self.model.backbone.named_parameters():
    print(name, torch.isnan(param.grad).any())

###### only noise to vs
layer3_custom_block1.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.bias tensor(True, device='cuda:0')
layer3_custom_block2.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.bias tensor(True, device='cuda:0')

###### both noise to vs and gradient hook to new_vt
layer3_custom_block1.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block1.fc_b.bias tensor(False, device='cuda:0')
layer3_custom_block2.fc1.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc2.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_w.bias tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.weight tensor(True, device='cuda:0')
layer3_custom_block2.fc_b.bias tensor(False, device='cuda:0')

So the nan gradient of new_vt comes from multiplication instead of std? ? The only reason I can see for multiplication to give nan gradient is the input is nan as well. But that doesn’t make sense since inserting noise to new_vt solves the problem. Is there another possible reason for multiplication to give nan gradient?
(P.S. I find this post https://github.com/pytorch/pytorch/issues/15288 2, but I’m not sure if it’s the same case)