I think you might see if track_running_stats in BatchNormXd (torch 0.4 and above) 105 does what you want.
You can probably copypaste the relevant code if you have an older version.
Best regards
Thomas