class_docstr: A generic `ModelHook` that ensures inputs and model weights are on the same device for the forward pass of the
  associated module, potentially offloading the weights after the forward pass.
  
  Args:
      execution_device (`torch.device`, *optional*):
          The device on which inputs and model weights should be placed before the forward pass.
      offload (`bool`, *optional*, defaults to `False`):
          Whether or not the weights should be offloaded after the forward pass.
      io_same_device (`bool`, *optional*, defaults to `False`):
          Whether or not the output should be placed on the same device as the input was.
      weights_map (`Mapping[str, torch.Tensor]`, *optional*):
          When the model weights are offloaded, a (potentially lazy) map from param names to the tensor values.
      offload_buffers (`bool`, *optional*, defaults to `False`):
          Whether or not to include the associated module's buffers when offloading.
      place_submodules (`bool`, *optional*, defaults to `False`):
          Whether to place the submodules on `execution_device` during the `init_hook` event.
code: def __init__(
      self,
      execution_device: Optional[Union[int, str, torch.device]] = None,
      offload: bool = False,
      io_same_device: bool = False,
      weights_map: Optional[Mapping] = None,
      offload_buffers: bool = False,
      place_submodules: bool = False,
      skip_keys: Optional[Union[str, List[str]]] = None,
  ):
      self.execution_device = execution_device
      self.offload = offload
      self.io_same_device = io_same_device
      self.weights_map = weights_map
      self.offload_buffers = offload_buffers
      self.place_submodules = place_submodules
      self.skip_keys = skip_keys
  
      # Will contain the input device when `io_same_device=True`.
      self.input_device = None
      self.param_original_devices = {}
      self.buffer_original_devices = {}
  
imports: import functools
  from typing import Dict, List, Mapping, Optional, Union
  import torch
  import torch.nn
  from .state import PartialState
  from .utils import PrefixedDataset, find_device, named_module_tensors, send_to_device, set_module_tensor_to_device
  from .utils.modeling import get_non_persistent_buffers
label_desc: Use image node_id 1 to create a new AlignDevicesHook object from inherited base classes: ModelHook with example: obj = AlignDevicesHook(execution_device, offload, io_same_device, weights_map, offload_buffers, place_submodules, skip_keys)