code: def forward(
      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
  
file_path: apex/apex/parallel/optimized_sync_batchnorm_kernel.py
imports: import torch
  from torch.autograd.function import Function
  import syncbn
  from apex.parallel import ReduceOp
label_desc: Use image node_id 1 for calling the SyncBatchnormFunction obj's underlying member method code with example usage: obj.forward(ctx, input, z, weight, bias, running_mean, running_variance, eps, track_running_stats, momentum, process_group, channel_last, fuse_relu) and returns: out