code: def with_temp_dir(
      func: Optional[Callable] = None,
  ) -> Optional[Callable]:
      """
      Wrapper to initialize temp directory for distributed checkpoint.
      """
      assert func is not None
  
      @wraps(func)
      def wrapper(
          self, *args: Tuple[object], **kwargs: Dict[str, Any]
      ) -> None:
          # Only create temp_dir when rank is 0
          if dist.get_rank() == 0:
              temp_dir = tempfile.mkdtemp()
              print(f"Using temp directory: {temp_dir}")
          else:
              temp_dir = ""
          object_list = [temp_dir]
  
          # Broadcast temp_dir to all the other ranks
          dist.broadcast_object_list(object_list)
          self.temp_dir = object_list[0]
  
          try:
              func(self, *args, **kwargs)
          finally:
              if dist.get_rank() == 0:
                  shutil.rmtree(self.temp_dir, ignore_errors=True)
  
      return wrapper
  
imports: import shutil
  import tempfile
  from functools import wraps
  from typing import Any, Callable, Dict, Optional, Tuple
  import torch
  import torch.distributed
  import torch.distributed.checkpoint
  import torch.nn
  from torch.distributed.checkpoint._fsspec_filesystem import FsspecReader, FsspecWriter
  from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType
  from torch.testing._internal.common_distributed import requires_nccl, skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import run_tests, TestCase
  from torch.testing._internal.distributed._shard.sharded_tensor import ShardedTensorTestBase, with_comms
label_desc: Use image node_id 1 for calling a global function with example usage: with_temp_dir(func) and returns: wrapper