code: def load_state_dict(checkpoint_file, device_map=None):
      """
      Load a checkpoint from a given file. If the checkpoint is in the safetensors format and a device map is passed, the
      weights can be fast-loaded directly on the GPU.
  
      Args:
          checkpoint_file (`str`): The path to the checkpoint to load.
          device_map (`Dict[str, Union[int, str, torch.device]]`, *optional*):
              A map that specifies where each submodule should go. It doesn't need to be refined to each parameter/buffer
              name, once a given module name is inside, every submodule of it will be sent to the same device.
      """
      if checkpoint_file.endswith(".safetensors"):
          with safe_open(checkpoint_file, framework="pt") as f:
              metadata = f.metadata()
              weight_names = f.keys()
  
          if metadata is None:
              logger.warn(
                  f"The safetensors archive passed at {checkpoint_file} does not contain metadata. "
                  "Make sure to save your model with the `save_pretrained` method. Defaulting to 'pt' metadata."
              )
              metadata = {"format": "pt"}
  
          if metadata.get("format") not in ["pt", "tf", "flax"]:
              raise OSError(
                  f"The safetensors archive passed at {checkpoint_file} does not contain the valid metadata. Make sure "
                  "you save your model with the `save_pretrained` method."
              )
          elif metadata["format"] != "pt":
              raise ValueError(
                  f"The checkpoint passed was saved with {metadata['format']}, we need a the pt format."
              )
          if device_map is None:
              return safe_load_file(checkpoint_file)
          else:
              # if we only have one device we can load everything directly
              if len(set(device_map.values())) == 1:
                  return safe_load_file(
                      checkpoint_file,
                      device=list(device_map.values())[0],
                  )
  
              devices = list(set(device_map.values()) - {"disk"})
              # cpu device should always exist as fallback option
              if "cpu" not in devices:
                  devices.append("cpu")
  
              # For each device, get the weights that go there
              device_weights = {device: [] for device in devices}
              for module_name, device in device_map.items():
                  if device in devices:
                      device_weights[device].extend(
                          [
                              k
                              for k in weight_names
                              if k == module_name
                              or k.startswith(module_name + ".")
                          ]
                      )
  
              # all weights that haven't defined a device should be loaded on CPU
              device_weights["cpu"].extend(
                  [
                      k
                      for k in weight_names
                      if k not in sum(device_weights.values(), [])
                  ]
              )
              tensors = {}
              if is_tqdm_available():
                  progress_bar = tqdm(
                      main_process_only=False,
                      total=sum(
                          [
                              len(device_weights[device])
                              for device in devices
                          ]
                      ),
                      unit="w",
                      smoothing=0,
                      leave=False,
                  )
              else:
                  progress_bar = None
              for device in devices:
                  with safe_open(
                      checkpoint_file, framework="pt", device=device
                  ) as f:
                      for key in device_weights[device]:
                          if progress_bar is not None:
                              progress_bar.set_postfix(
                                  dev=device, refresh=False
                              )
                              progress_bar.set_description(key)
                          tensors[key] = f.get_tensor(key)
                          if progress_bar is not None:
                              progress_bar.update()
              if progress_bar is not None:
                  progress_bar.close()
  
              return tensors
      else:
          return torch.load(
              checkpoint_file, map_location=torch.device("cpu")
          )
  
imports: import contextlib
  import gc
  import inspect
  import json
  import logging
  import os
  import re
  import shutil
  import tempfile
  from collections import OrderedDict, defaultdict
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  import torch.nn
  from ..state import AcceleratorState
  from .constants import SAFE_WEIGHTS_NAME, WEIGHTS_NAME
  from .dataclasses import AutocastKwargs, CustomDtype, DistributedType
  from .imports import is_mps_available, is_npu_available, is_xpu_available
  from .offload import load_offloaded_weight, offload_weight, save_offload_index
  from .tqdm import is_tqdm_available, tqdm
  from safetensors import safe_open
  from safetensors.torch import load_file
label_desc: Use image node_id 24 for calling a global function with example usage: load_state_dict(checkpoint_file, device_map) and returns: torch, safe_load_file, tensors, safe_load_file