Well, I don’t know if the hook keeps the anomaly mode from complaining and while I don’t fully understand what you are doing, but it would seem that std still is the critical bit.
One (slow, but maybe OK for debugging) alternative to the gradient hook is to insert some noise new_vt = new_vt + 1e-6 * torch.randn_like(new_vt) in the forward.
Best regards
Thomas