def METHOD_NAME(model):
    pdb.set_trace()
    print("Not sure if using this is a good idea")
    modules = model._modules
    for m in modules.keys():
        module = modules[m]
        if isinstance(module, nn.BatchNorm2d) or isinstance(module, nn.BatchNorm1d):
            feature_number = module.num_features
            model._modules[m] = nn.GroupNorm(32, feature_number)
        elif isinstance(module, nn.BatchNorm3d):
            raise Exception("Not implemented")
        elif isinstance(module, nn.Module):
            model._modules[m] = METHOD_NAME(module)
    return model