How to normalize feature using running mean and running var meanwhile keep update them?

In my mind, the eval mode will normalize feature maps with the recorded running mean and running var. But in eval mode, they cannot be updated.

In train mode, the running mean and running var can be updated, but the BN layer doesn’t use them to normalize feature.

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.