code: def save(
      obj,
      f,
      save_on_each_node: bool = False,
      safe_serialization: bool = False,
  ):
      """
      Save the data to disk. Use in place of `torch.save()`.
  
      Args:
          obj:
              The data to save
          f:
              The file (or file-like object) to use to save the data
          save_on_each_node (`bool`, *optional*, defaults to `False`):
              Whether to only save on the global main process
          safe_serialization (`bool`, *optional*, defaults to `False`):
              Whether to save `obj` using `safetensors` or the traditional PyTorch way (that uses `pickle`).
      """
      # Check if it's a model and remove duplicates
      if safe_serialization:
          save_func = partial(safe_save_file, metadata={"format": "pt"})
          if isinstance(obj, OrderedDict):
              obj = clean_state_dict_for_safetensors(obj)
      else:
          save_func = torch.save
  
      if PartialState().distributed_type == DistributedType.TPU:
          xm.save(obj, f)
      elif PartialState().is_main_process and not save_on_each_node:
          save_func(obj, f)
      elif PartialState().is_local_main_process and save_on_each_node:
          save_func(obj, f)
  
imports: import collections
  import os
  import platform
  import re
  import socket
  from contextlib import contextmanager
  from functools import partial
  from types import MethodType
  from typing import OrderedDict
  import torch
  from packaging.version import Version
  from safetensors.torch import save_file
  from ..commands.config.default import write_basic_config
  from ..logging import get_logger
  from ..state import PartialState
  from .constants import FSDP_PYTORCH_VERSION
  from .dataclasses import DistributedType
  from .imports import is_deepspeed_available, is_torch_distributed_available, is_tpu_available
  from .modeling import id_tensor_storage
  from .transformer_engine import convert_model
  from .versions import is_torch_version
label_desc: Use image node_id 5 for calling a global function with example usage: save(obj, f, save_on_each_node, safe_serialization) without return types