Shanyan_Guan:

How to update running mean and var and then use them to normalize the input features?


Do you want to calculate the batch stats, update the running stats, and use the running stats to normalize the data?
If so, I think the cleanest way would be to implement a custom batchnorm layer by changing the logic in the current implementation.