Hi, i know it’s way to late.
But for everyone, who is trying to replace batch norm with group norm for more complex models like you wanted to do, you can use the following hacky approach :
def batch_norm_to_group_norm(layer):
    """Iterates over a whole model (or layer of a model) and replaces every batch norm 2D with a group norm

    Args:
        layer: model or one layer of a model like resnet34.layer1 or Sequential(), ...
    """
    for name, module in layer.named_modules():
        if name:
            try:
                # name might be something like: model.layer1.sequential.0.conv1 --> this wont work. Except this case
                sub_layer = getattr(layer, name)
                if isinstance(sub_layer, torch.nn.BatchNorm2d):
                    num_channels = sub_layer.num_features
                    # first level of current layer or model contains a batch norm --> replacing.
                    layer._modules[name] = torch.nn.GroupNorm(constants.GROUP_NORM_LOOKUP[num_channels], num_channels)
            except AttributeError:
                # go deeper: set name to layer1, getattr will return layer1 --> call this func again
                name = name.split('.')[0]
                sub_layer = getattr(layer, name)
                sub_layer = batch_norm_to_group_norm(sub_layer)
                layer.__setattr__(name=name, value=sub_layer)
    return layer

And the GROUP_NORM_LOOKUP is build as follows:
# group norm paper stated, that 16 channels per group are best
# group norm paper stated, that 8 or 32 groups are best
# channels per group = num_features / num_group
# channels per group have less influence when being >= 8 so here i try to set optimal values.
# Paper values:
# groups (G)
# 64   32   16   8    4    2    1 (=LN)
# 24.6 24.1 24.6 24.4 24.6 24.7 25.3  # validation error in %
# 0.5  -    0.5  0.3  0.5  0.6  1.2 
# channels per group
# 64   32   16   8    4    2    1 (=IN)
# 24.4 24.5 24.2 24.3 24.8 25.6 28.4  # validation error in %
# 0.2  0.3  -    0.1  0.6  1.4  4.2

# num_channels: num_groups
GROUP_NORM_LOOKUP = {
    16: 2,  # -> channels per group: 8
    32: 4,  # -> channels per group: 8
    64: 8,  # -> channels per group: 8
    128: 8,  # -> channels per group: 16
    256: 16,  # -> channels per group: 16
    512: 32,  # -> channels per group: 16
    1024: 32,  # -> channels per group: 32
    2048: 32,  # -> channels per group: 64
}

Of course you can use a different approach for setting the number of groups.