Replace BN layers in model

Hi,

I am using two models, DenseNet and EfficientNet (B7) and would like to replace batch norm with group norm (https://arxiv.org/pdf/1803.08494.pdf) or weight standardization (https://arxiv.org/pdf/1903.10520.pdf).

I have found several implementations of these norm techniques, however, I could not find how to edit the model (both of them are from libraries and I cannot edit the source code directly). Is there some way, how to replace all BN layers by GN od WS in PyTorch?

Thanks
Best regards