code: def save_state_dict(
      state_dict: dict, checkpoint_file_path: str, use_safetensors: bool
  ) -> None:
      """
      Save state dict to checkpoint.
  
      Args:
          state_dict (dict): state dict.
          checkpoint_file_path (str): path to the checkpoint file.
          use_safetensors (bool): whether to use safetensors to save the checkpoint.
      """
      # Move all tensors in the state_dict to CPU before saving to avoid serialization issues
      state_dict_cpu = tree_map(
          lambda x: x.cpu() if torch.is_tensor(x) else x, state_dict
      )
  
      if use_safetensors:
          assert (
              is_safetensors_available()
          ), "safetensors is not available."
          assert checkpoint_file_path.endswith(
              ".safetensors"
          ), "safetensors only supports .safetensors suffix for checkpoint file."
          from safetensors.torch import save_file as safe_save_file
  
          safe_save_file(
              state_dict_cpu,
              checkpoint_file_path,
              metadata={"format": "pt"},
          )
      else:
          torch.save(state_dict_cpu, checkpoint_file_path)
  
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 10 for calling a global function with example usage: save_state_dict(state_dict, checkpoint_file_path, use_safetensors) without return types