code: def pg_test(world_size=torch.cuda.device_count()):
      def decorator(func):
          @functools.wraps(func)
          def wrapper(*args, **kwargs):
              mp.spawn(
                  pg_worker,
                  args=(
                      world_size,
                      tempfile.mkstemp()[1],
                      func,
                      *kwargs.values(),
                  ),
                  nprocs=world_size,
              )
  
          globals()["test_" + func.__name__] = wrapper
          return func
  
      return decorator
  
imports: import functools
  import tempfile
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  from torch.nn.parallel import DistributedDataParallel
  from fairscale.experimental.nn import SyncBatchNorm
  from fairscale.nn.checkpoint import checkpoint_wrapper
label_desc: Use image node_id 2 for calling a global function with example usage: pg_test(world_size) and returns: decorator, func