code: def load_state_dict(checkpoint_file_path: Path):
      """
      Load state dict from checkpoint.
  
      Args:
          checkpoint_file_path (Path): path to the checkpoint file.
  
      Returns:
          dict: state dict.
      """
  
      assert not is_dtensor_checkpoint(
          checkpoint_file_path
      ), f"Cannot load state dict from dtensor checkpoint {checkpoint_file_path}, you should convert the distributed tensors to gathered tensors with our CLI offline."
  
      if is_safetensor_checkpoint(checkpoint_file_path):
          assert (
              is_safetensors_available()
          ), f"Cannot load state dict from safetensor checkpoint {checkpoint_file_path}, because safetensors is not available. Please install safetensors first with pip install safetensors."
          # load with safetensors
          from safetensors import safe_open
  
          state_dict = {}
          with safe_open(
              checkpoint_file_path, framework="pt", device="cpu"
          ) as f:
              for k in f.keys():
                  state_dict[k] = f.get_tensor(k)
          return state_dict
  
      else:
          # load with torch
          return torch.load(
              checkpoint_file_path, map_location=torch.device("cpu")
          )
  
imports: import os
  import re
  from collections import abc
  from collections import defaultdict
  from itertools import chain
  from pathlib import Path
  from typing import Iterator, List, Mapping, Optional, OrderedDict, Tuple
  import torch
  import torch.nn
  from packaging.version import Version
  from torch.optim import Optimizer
  from torch.utils._pytree import tree_map
  from colossalai.tensor.d_tensor import is_customized_distributed_tensor, is_distributed_tensor, to_global, to_global_for_customized_distributed_tensor
label_desc: Use image node_id 24 for calling a global function with example usage: load_state_dict(checkpoint_file_path) and returns: state_dict, torch