code: def forward(
      self,
      input,
      weight,
      bias,
      running_mean,
      running_var,
      eps,
      momentum,
  ):
      input = input.contiguous()
  
      size = input.numel() // input.size(1)
      count = torch.tensor([size])
  
      # calculate mean/invstd for input.
      mean, invstd = torch.batch_norm_stats(input, eps)
  
      count_handle = allgather_async(
          count.unsqueeze(0), name="sync_batch_norm.count"
      )
      mean_handle = allgather_async(
          mean.unsqueeze(0), name="sync_batch_norm.mean"
      )
      invstd_handle = allgather_async(
          invstd.unsqueeze(0), name="sync_batch_norm.invstd"
      )
  
      # wait on the async communication to finish
      count_all = synchronize(count_handle)
      mean_all = synchronize(mean_handle)
      invstd_all = synchronize(invstd_handle)
  
      if _SYNC_BN_V3:
          counts_for_bngswc = (
              count_all.view(-1).float().to(input.device)
          )
      else:
          # backwards compatibility
          counts_for_bngswc = count_all.view(-1).tolist()
  
      # calculate global mean & invstd
      mean, invstd = torch.batch_norm_gather_stats_with_counts(
          input,
          mean_all,
          invstd_all,
          running_mean,
          running_var,
          momentum,
          eps,
          counts_for_bngswc,
      )
  
      self.save_for_backward(input, weight, mean, invstd, count_all)
  
      # apply element-wise normalization
      return torch.batch_norm_elemt(
          input, weight, bias, mean, invstd, eps
      )
  
file_path: horovod/horovod/torch/sync_batch_norm.py
imports: from horovod.torch.mpi_ops import allgather_async, allreduce_async, Sum, size, synchronize
  from packaging import version
  import torch
  from torch.autograd.function import Function
  import torch.nn.functional
  from torch.nn.modules.batchnorm import _BatchNorm
label_desc: Use image node_id 1 for calling the _SyncBatchNorm obj's underlying member method code with example usage: obj.forward(input, weight, bias, running_mean, running_var, eps, momentum) and returns: torch