Solved by ptrblck in post #4 
              
              
                What about defining a method which returns the desired layer based on your keyword: 
def get_norm_layer(out_channels, layer_type='BatchNorm'):
    if layer_type == 'BatchNorm':
        return nn.BatchNorm2d(out_channels)
    elif layer_type == 'GroupNorm':
        return nn.GroupNorm(16, out_channel…