ConvModule with layernorm

[Stephanie-Shen324]

Hi, is it possible to use layernorm in convModule?

nn.layernorm expects a tensor with channel last (N, H, W, C) where normalized shape is (C). However the convModule's input/output are channel-first (N, C, H, W). hence running convModule with norm_cfg(type='LN') will give the error of "Given normalized_shape=[256], expected input with shape [*, 256], but got input of size[batch_size, 256, H, W]"

However, I did get it running with no errors with norm_cfg(type='GN', num_groups=1). And from the following image, group norm with 1 group should be equivalent to layernorm?

Please let me know if i made a mistake somewhere and how can I make it work for norm_cfg(type='LN') in ConvModule.