code: def prepare_fsdp(
      module: torch.nn.Module,
      device: torch.device,
      strategy: Optional[FSDPStrategy] = None,
  ) -> FSDP:
      """
      Utility to move a module to device and wrap in `FullyShardedDataParallel <https://pytorch.org/docs/stable/fsdp.html>`_.
  
      Args:
          module: module to be wrapped in FSDP
          device: device to which module will be moved
          strategy: an instance of :class:`~torchtnt.utils.prepare_module.FSDPStrategy` which defines the settings of FSDP APIs
  
      Examples::
          strategy = FSDPStrategy(limit_all_gathers=True)
          module = nn.Linear(1, 1)
          device = torch.device("cuda")
          fsdp_module = prepare_fsdp(module, device, strategy)
      """
      if not is_torch_version_geq_1_12():
          raise RuntimeError(
              "Please install PyTorch 1.12 or higher to use FSDP: https://pytorch.org/get-started/locally/"
          )
      strategy = strategy if strategy is not None else FSDPStrategy()
  
      # we use __dict__ and not asdict() here because asdict() is recursively applied on nested objects
      params_dict = strategy.__dict__.copy()
  
      # extract params to set state dict type
      state_dict_type = params_dict.pop("state_dict_type")
      state_dict_config = params_dict.pop("state_dict_config")
      optim_state_dict_config = params_dict.pop(
          "optim_state_dict_config"
      )
  
      # wrap module in FSDP
      module = FSDP(
          module,
          device_id=device,
          **params_dict,
      )
  
      if state_dict_type:
          FSDP.set_state_dict_type(
              module,
              state_dict_type,
              state_dict_config,
              optim_state_dict_config,
          )
      return module
  
imports: from dataclasses import asdict, dataclass
  from functools import partial
  from typing import Any, Callable, Dict, Iterable, Optional, Union
  import torch
  import torch.distributed
  from torch.distributed import ProcessGroup
  from torch.distributed.fsdp import FullyShardedDataParallel, StateDictType
  from torch.distributed.fsdp.api import OptimStateDictConfig, StateDictConfig
  from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch, CPUOffload, MixedPrecision, ShardingStrategy
  from torch.nn.parallel import DistributedDataParallel
  from torchtnt.utils.rank_zero_log import rank_zero_warn
  from torchtnt.utils.version import is_torch_version_geq_1_12, is_torch_version_geq_2_0
label_desc: Use image node_id 2 for calling a global function with example usage: prepare_fsdp(module, device, strategy) and returns: module