code: def init_on_device(
      device: torch.device, include_buffers: bool = None
  ):
      """
      A context manager under which models are initialized with all parameters on the specified device.
  
      Args:
          device (`torch.device`):
              Device to initialize all parameters on.
          include_buffers (`bool`, *optional*):
              Whether or not to also put all buffers on the meta device while initializing.
  
      Example:
  
      ```python
      import torch.nn as nn
      from accelerate import init_on_device
  
      with init_on_device(device=torch.device("cuda")):
          tst = nn.Liner(100, 100)  # on `cuda` device
      ```
      """
      if include_buffers is None:
          include_buffers = parse_flag_from_env(
              "ACCELERATE_INIT_INCLUDE_BUFFERS", False
          )
  
      # TODO(shingjan): remove the torch version check once older versions are deprecated
      if is_torch_version(">=", "2.0") and include_buffers:
          with device:
              yield
          return
  
      old_register_parameter = nn.Module.register_parameter
      if include_buffers:
          old_register_buffer = nn.Module.register_buffer
  
      def register_empty_parameter(module, name, param):
          old_register_parameter(module, name, param)
          if param is not None:
              param_cls = type(module._parameters[name])
              kwargs = module._parameters[name].__dict__
              module._parameters[name] = param_cls(
                  module._parameters[name].to(device), **kwargs
              )
  
      def register_empty_buffer(module, name, buffer, persistent=True):
          old_register_buffer(
              module, name, buffer, persistent=persistent
          )
          if buffer is not None:
              module._buffers[name] = module._buffers[name].to(device)
  
      # Patch tensor creation
      if include_buffers:
          tensor_constructors_to_patch = {
              torch_function_name: getattr(torch, torch_function_name)
              for torch_function_name in [
                  "empty",
                  "zeros",
                  "ones",
                  "full",
              ]
          }
      else:
          tensor_constructors_to_patch = {}
  
      def patch_tensor_constructor(fn):
          def wrapper(*args, **kwargs):
              kwargs["device"] = device
              return fn(*args, **kwargs)
  
          return wrapper
  
      try:
          nn.Module.register_parameter = register_empty_parameter
          if include_buffers:
              nn.Module.register_buffer = register_empty_buffer
          for (
              torch_function_name
          ) in tensor_constructors_to_patch.keys():
              setattr(
                  torch,
                  torch_function_name,
                  patch_tensor_constructor(
                      getattr(torch, torch_function_name)
                  ),
              )
          yield
      finally:
          nn.Module.register_parameter = old_register_parameter
          if include_buffers:
              nn.Module.register_buffer = old_register_buffer
          for (
              torch_function_name,
              old_torch_function,
          ) in tensor_constructors_to_patch.items():
              setattr(torch, torch_function_name, old_torch_function)
  
file_path: accelerate/src/accelerate/big_modeling.py
imports: import logging
  import os
  from contextlib import contextmanager
  from functools import wraps
  from typing import Dict, List, Optional, Union
  import torch
  import torch.nn
  from .hooks import AlignDevicesHook, CpuOffload, UserCpuOffloadHook, add_hook_to_module, attach_align_device_hook, attach_align_device_hook_on_blocks
  from .utils import OffloadedWeightsLoader, check_device_map, extract_submodules_state_dict, find_tied_parameters, get_balanced_memory, infer_auto_device_map, is_npu_available, is_torch_version, load_checkpoint_in_model, offload_state_dict, parse_flag_from_env, retie_parameters
label_desc: Use image node_id 2 for calling a global function with example usage: init_on_device(device, include_buffers) and returns: wrapper, fn