code: def prepare_module(
      module: torch.nn.Module,
      device: torch.device,
      *,
      strategy: Optional[Union[Strategy, str]] = None,
      torch_compile_params: Optional[TorchCompileParams] = None,
      activation_checkpoint_params: Optional[
          ActivationCheckpointParams
      ] = None,
  ) -> torch.nn.Module:
      """
      Utility to move a module to device, set up parallelism, activation checkpointing and compile.
  
      Args:
          module: module to be used.
          device: device to which module will be moved.
          strategy: the data parallelization strategy to be used. if a string, must be one of ``ddp``, ``fsdp``, or ``noop``.
          torch_compile_params: params for Torch compile https://pytorch.org/docs/stable/generated/torch.compile.html.
          activation_checkpoint_params: params for enabling activation checkpointing.
      """
  
      if strategy:
          if isinstance(strategy, str):
              strategy = convert_str_to_strategy(strategy)
          if isinstance(strategy, DDPStrategy):
              if torch_compile_params and strategy.static_graph is True:
                  # https://dev-discuss.pytorch.org/t/torchdynamo-update-9-making-ddp-work-with-torchdynamo/860
                  raise RuntimeError(
                      "Torch compile requires DDPStrategy's static_graph to be False"
                  )
              module = prepare_ddp(module, device, strategy)
          elif isinstance(strategy, FSDPStrategy):
              if (
                  torch_compile_params
                  and strategy.use_orig_params is False
              ):
                  # as stated here https://pytorch.org/get-started/pytorch-2.0/
                  raise RuntimeError(
                      "Torch compile requires FSDPStrategy's use_orig_params to be True, since AOTAutograd needs to be aware of the original parameters"
                  )
              module = prepare_fsdp(module, device, strategy)
      else:
          module = module.to(device)
  
      if activation_checkpoint_params:
          if not is_torch_version_geq_2_0():
              raise RuntimeError(
                  "Activation checkpointing requires torch>=2.0"
              )
          checkpoint_impl = activation_checkpoint_params.checkpoint_impl
          check_fn = activation_checkpoint_params.check_fn
          auto_wrap_policy = (
              activation_checkpoint_params.auto_wrap_policy
          )
          custom_checkpoint_wrapper = partial(
              checkpoint_wrapper,
              checkpoint_impl=checkpoint_impl,
          )
          apply_activation_checkpointing(
              module,
              checkpoint_wrapper_fn=custom_checkpoint_wrapper,
              check_fn=check_fn,
              auto_wrap_policy=auto_wrap_policy,
          )
  
      if torch_compile_params:
          try:
              # use in-place compile to avoid altering the state_dict keys
              module.compile(**asdict(torch_compile_params))
          except AttributeError:
              rank_zero_warn(
                  "Please install PyTorch nightlies to use in-place compile to avoid altering the state_dict keys when checkpointing."
              )
              torch.compile(module, **asdict(torch_compile_params))
  
      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 4 for calling a global function with example usage: prepare_module(module, device) and returns: module