def METHOD_NAME(ctx, input, z, weight, bias, running_mean, running_variance, eps, track_running_stats = True, momentum = 1.0, process_group = None, channel_last = False, fuse_relu = False):
    input = input.contiguous()
    world_size = 0
    mean = None
    var_biased = None
    inv_std = None
    var = None
    out = None
    count = None
    if track_running_stats:
        if channel_last:
            count = int(input.numel()/input.size(-1))
            mean, var_biased = syncbn.welford_mean_var_c_last(input)
            num_channels = input.size(-1)
        else:
            count = int(input.numel()/input.size(1))
            mean, var_biased = syncbn.welford_mean_var(input)
            num_channels = input.size(1)
        if torch.distributed.is_initialized():
            if not process_group:
                process_group = torch.distributed.group.WORLD
            device = mean.device
            world_size = torch.distributed.get_world_size(process_group)
            count_t = torch.empty(1, dtype=mean.dtype, device=mean.device).fill_(count)
            combined = torch.cat([mean.view(-1), var_biased.view(-1), count_t], dim=0)
            combined_list = [torch.empty_like(combined) for k in range(world_size)]
            torch.distributed.all_gather(combined_list, combined, process_group)
            combined = torch.stack(combined_list, dim=0)
            mean_all, invstd_all, count_all = torch.split(combined, num_channels, dim=1)
            count_all = count_all.view(-1)
            mean, var, inv_std = syncbn.welford_parallel(mean_all, invstd_all, count_all.to(torch.int32), eps)
        else:
            device = mean.device
            count_all = torch.cuda.IntTensor([count], device=device)
            inv_std = 1.0 / torch.sqrt(var_biased + eps)
            var = var_biased * (count) / (count-1)
        if count == 1 and world_size < 2:
            raise ValueError('Expected more than 1 value per channel when training, got input size{}'.format(input.size()))
        r_m_inc = mean if running_mean.dtype != torch.float16 else mean.half()
        r_v_inc = var if running_variance.dtype != torch.float16 else var.half()
        running_mean.data = running_mean.data * (1-momentum) + momentum*r_m_inc
        running_variance.data = running_variance.data * (1-momentum) + momentum*r_v_inc
    else:
        mean = running_mean.data
        inv_std = 1.0 / torch.sqrt(running_variance.data + eps)
    ctx.save_for_backward(input, weight, mean, inv_std, z, bias, count_all.to(torch.int32))
    ctx.process_group = process_group
    ctx.channel_last = channel_last
    ctx.world_size = world_size
    ctx.fuse_relu = fuse_relu
    if channel_last:
        out = syncbn.batchnorm_forward_c_last(input, z, mean, inv_std, weight, bias, fuse_relu)
    else:
        out = syncbn.batchnorm_forward(input, mean, inv_std, weight, bias)
    return out