code: def with_comms(func=None):
      if func is None:
          return partial(
              with_comms,
          )
  
      @wraps(func)
      def wrapper(self, *args, **kwargs):
          if (
              BACKEND == dist.Backend.NCCL
              and torch.cuda.device_count() < self.world_size
          ):
              sys.exit(
                  TEST_SKIPS[f"multi-gpu-{self.world_size}"].exit_code
              )
          self.dist_init()
          func(self)
          self.destroy_comms()
  
      return wrapper
  
imports: import os
  import sys
  from functools import wraps, partial
  import torch
  import torch.distributed
  from torch.testing._internal.common_distributed import MultiProcessTestCase, TEST_SKIPS
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 1 for calling a global function with example usage: with_comms(func) and returns: wrapper, partial