code: def distribute_bn(model, world_size, reduce=False):
      # ensure every node has the same running bn stats
      for bn_name, bn_buf in unwrap_model(model).named_buffers(
          recurse=True
      ):
          if ("running_mean" in bn_name) or ("running_var" in bn_name):
              if reduce:
                  # average bn stats across whole group
                  torch.distributed.all_reduce(
                      bn_buf, op=dist.ReduceOp.SUM
                  )
                  bn_buf /= float(world_size)
              else:
                  # broadcast bn stats from rank 0 to whole group
                  torch.distributed.broadcast(bn_buf, 0)
  
file_path: pytorch-image-models/timm/utils/distributed.py
imports: import os
  import torch
  from torch import distributed
  from .model import unwrap_model
label_desc: Use image node_id 2 for calling a global function with example usage: distribute_bn(model, world_size, reduce) without return types