class_docstr: A utility that implements methods to accelerate PyTorch training.
  
  Arguments:
      amp: If true, perform training with automatic mixed precision.
          Otherwise, use full precision.
class_docstr_tok: ["A","utility","that","implements","methods","to","accelerate","PyTorch","training",".","Arguments",":","amp",":","If","true",",","perform","training","with","automatic","mixed","precision",".","Otherwise",",","use","full","precision","."]
docstr: Prepares the model for distributed execution.
  
  This allows you to use the same exact code regardless of number of
  workers or the device type being used (CPU, GPU).
  
  Args:
      model (torch.nn.Module): A torch model to prepare.
      move_to_device: Whether to move the model to the correct
          device. If set to False, the model needs to manually be moved
          to the correct device.
      parallel_strategy ("ddp", "fsdp", or None): Whether to wrap models
          in ``DistributedDataParallel``, ``FullyShardedDataParallel`` (
          Experimental), or neither.
      parallel_strategy_kwargs (Dict[str, Any]): Args to pass into
          ``DistributedDataParallel`` or ``FullyShardedDataParallel``
          initialization if ``parallel_strategy`` is set to "ddp"
          or "fsdp", respectively.
docstr_tok: ["Prepares","the","model","for","distributed","execution",".","This","allows","you","to","use","the","same","exact","code","regardless","of","number","of","workers","or","the","device","type","being","used","(","CPU",",","GPU",")",".","Args",":","model","(","torch.nn.Module",")",":","A","torch","model","to","prepare",".","move_to_device",":","Whether","to","move","the","model","to","the","correct","device",".","If","set","to","False",",","the","model","needs","to","manually","be","moved","to","the","correct","device",".","parallel_strategy","(","``","ddp","''",",","``","fsdp","''",",","or","None",")",":","Whether","to","wrap","models","in","``","DistributedDataParallel","``",",","``","FullyShardedDataParallel","``","(","Experimental",")",",","or","neither",".","parallel_strategy_kwargs","(","Dict","[","str",",","Any","]",")",":","Args","to","pass","into","``","DistributedDataParallel","``","or","``","FullyShardedDataParallel","``","initialization","if","``","parallel_strategy","``","is","set","to","``","ddp","''","or","``","fsdp","''",",","respectively","."]
code: def prepare_model(
      self,
      model: torch.nn.Module,
      move_to_device: bool = True,
      parallel_strategy: Optional[str] = "ddp",
      parallel_strategy_kwargs: Optional[Dict[str, Any]] = None,
  ) -> torch.nn.Module:
      """Prepares the model for distributed execution.
  
      This allows you to use the same exact code regardless of number of
      workers or the device type being used (CPU, GPU).
  
      Args:
          model (torch.nn.Module): A torch model to prepare.
          move_to_device: Whether to move the model to the correct
              device. If set to False, the model needs to manually be moved
              to the correct device.
          parallel_strategy ("ddp", "fsdp", or None): Whether to wrap models
              in ``DistributedDataParallel``, ``FullyShardedDataParallel`` (
              Experimental), or neither.
          parallel_strategy_kwargs (Dict[str, Any]): Args to pass into
              ``DistributedDataParallel`` or ``FullyShardedDataParallel``
              initialization if ``parallel_strategy`` is set to "ddp"
              or "fsdp", respectively.
      """
      parallel_strategy_kwargs = parallel_strategy_kwargs or {}
  
      rank = session.get_local_rank()
  
      if isinstance(move_to_device, torch.device):
          device = move_to_device
      else:
          device = get_device()
          if isinstance(device, list):
              device = device[0]
  
      if torch.cuda.is_available():
          torch.cuda.set_device(device)
  
      if move_to_device:
          if rank == 0:
              logger.info(f"Moving model to device: {device}")
          else:
              logger.debug(f"Moving model to device: {device}")
          model = model.to(device)
  
      def model_get_state(self):
          # `__getstate__` is an special method that informs pickle which attributes
          # to serialize. This custom implementation ensures that the wrapped forward
          # method and custom `__getstate__` method aren't serialized.
          if hasattr(self, "_original_get_state"):
              state = self._original_get_state()
              state["__getstate__"] = state["_original_get_state"]
              del state["_original_get_state"]
          else:
              # If model does not have a `__getstate__` already defined, use default
              # implementation.
              state = self.__dict__.copy()
              del state["__getstate__"]
          state["forward"] = state["_unwrapped_forward"]
          del state["_unwrapped_forward"]
  
          return state
  
      if self.amp_is_enabled:
          # Pickle cannot serialize the wrapped forward method. As a workaround,
          # define a custom `__getstate__` method that unwraps the forward method.
          model._unwrapped_forward = model.forward
          model.forward = autocast()(model.forward)
  
          # TODO(amogkam): Replace below logic with a generic "unpack model" method.
          # Replacing the `model.forward` method makes the model no longer
          # serializable. When serializing the model, we have to override the
          # `__getstate__` method to set back the original forward method.
          if hasattr(model, "__getstate__"):
              model._original_get_state = model.__getstate__
          # `__getstate__` must be a bound method rather than an callable attribute.
          # See https://stackoverflow.com/questions/972/adding-a-method-to-an-existing-object-instance.  # noqa: E501
          model.__getstate__ = types.MethodType(model_get_state, model)
  
      world_size = session.get_world_size()
  
      if parallel_strategy and world_size > 1:
          if parallel_strategy == "ddp":
              DataParallel = DistributedDataParallel
              if torch.cuda.is_available():
                  parallel_strategy_kwargs = {
                      "device_ids": [device],
                      "output_device": device,
                      **parallel_strategy_kwargs,
                  }
          else:
              if not torch.cuda.is_available():
                  raise RuntimeError(
                      "FSDP is only available with GPU-enabled "
                      "training. Set "
                      "`use_gpu=True` in your Trainer to train with "
                      "GPUs."
                  )
              DataParallel = FullyShardedDataParallel
          if rank == 0:
              logger.info(
                  f"Wrapping provided model in {DataParallel.__name__}."
              )
          else:
              logger.debug(
                  f"Wrapping provided model in {DataParallel.__name__}."
              )
          model = DataParallel(model, **parallel_strategy_kwargs)
  
      return model
  
imports: import collections
  import logging
  import os
  import random
  import types
  from typing import Any, Callable, Dict, List, Optional, Union
  import numpy
  import torch
  from packaging.version import Version
  from torch.cuda.amp import GradScaler, autocast
  from torch.nn.parallel import DistributedDataParallel
  from torch.optim import Optimizer
  from torch.utils.data import DataLoader, DistributedSampler, IterableDataset, RandomSampler, SequentialSampler
  from ray._private.usage.usage_lib import TagKey, record_extra_usage_tag
  from ray.train._internal import session
  from ray.train._internal.accelerator import Accelerator
  from ray.train._internal.session import get_accelerator, set_accelerator
  from ray.util.annotations import Deprecated, PublicAPI
label_desc: Use image node_id 2 for calling the _TorchAccelerator obj's underlying member method code with example usage: obj.prepare_model(model, move_to_device, parallel_strategy, parallel_strategy_kwargs) and returns: model, state