Humm
You can use the register_hook() 11 function to have a function called when gradient is computed.
I would use that to check exactly where are the nans appearing here. Is it possible that std returns nan gradients for certain inputs?