code: def benchmark_bn(rank, world_size, init_file, bn_cls):
      dist.init_process_group(
          dist.Backend.NCCL,
          init_method="file://" + init_file,
          rank=rank,
          world_size=world_size,
      )
      x = torch.randn(50, 2048, 7, 7).to(rank)
      bn = bn_cls(2048).to(rank)
      bn = DDP(bn, device_ids=[rank])
      # Warmup
      for i in range(50):
          with torch.no_grad():
              x = bn(x)
      torch.cuda.synchronize(rank)
      t0 = time.time()
      for i in range(100):
          with torch.no_grad():
              x = bn(x)
      torch.cuda.synchronize(rank)
      t1 = time.time()
      print("Elapsed time is ", t1 - t0)
  
imports: import tempfile
  import time
  import torch
  import torch.distributed
  import torch.multiprocessing
  from torch.nn.parallel import DistributedDataParallel
  import fairscale.experimental.nn
label_desc: Use image node_id 1 for calling a global function with example usage: benchmark_bn(rank, world_size, init_file, bn_cls) without return types