How to realize an arbitrary dimension batchnorm layer?

Hi, I want an arbitrary batchnorm layer for my CNN-based progressor. could some one tell me how to realize it ?
Here is an example of an anbitrary conv layer: https://github.com/pvjosue/pytorch_convNd. I want to apply the batchnorm6d after conv6d.
thank you.

Cross post from here with a potential approach.