Could edit the batchnorm part in torch/nn/functional.py work for sync bn?

  
      github.com
  
  
    pytorch/pytorch/blob/91d28026f8e8386d45dfb7ddab84c906479dead1/torch/nn/modules/batchnorm.py#L450 14

        else:  # use exponential moving average
            exponential_average_factor = self.momentum


    world_size = 1
    process_group = torch.distributed.group.WORLD
    if self.process_group:
        process_group = self.process_group
    world_size = torch.distributed.get_world_size(process_group)


    # fallback to framework BN when synchronization is not necessary
    if world_size == 1 or (not self.training and self.track_running_stats):
        return F.batch_norm(
            input, self.running_mean, self.running_var, self.weight, self.bias,
            self.training or not self.track_running_stats,
            exponential_average_factor, self.eps)
    else:
        return sync_batch_norm.apply(
            input, self.weight, self.bias, self.running_mean, self.running_var,
            self.eps, exponential_average_factor, process_group, world_size)


@classmethod



  
  
    
    
  
  


The batch norm in torch.nn.functional is used just for evaluation. I think editing this would do nothing to sync batch norm. How do you edit the file to make sync bn work normally?