code: def load_shard_state_dict(
      checkpoint_file: Path, use_safetensors: bool = False
  ):
      """
      load shard state dict into model
      """
      if (
          use_safetensors
          and not checkpoint_file.suffix == ".safetensors"
      ):
          raise Exception(
              "load the model using `safetensors`, but no file endwith .safetensors"
          )
      if use_safetensors:
          from safetensors.torch import load_file as safe_load_file
          from safetensors.torch import safe_open
  
          with safe_open(checkpoint_file, framework="pt") as f:
              metadata = f.metadata()
          if metadata["format"] != "pt":
              raise NotImplementedError(
                  f"Conversion from a {metadata['format']} safetensors archive to PyTorch is not implemented yet."
              )
          return safe_load_file(checkpoint_file)
      else:
          return torch.load(
              checkpoint_file, 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 18 for calling a global function with example usage: load_shard_state_dict(checkpoint_file, use_safetensors) and returns: safe_load_file, torch