code: def load_state_dict(checkpoint_path, use_ema=False):
      if checkpoint_path and os.path.isfile(checkpoint_path):
          checkpoint = torch.load(checkpoint_path, map_location="cpu")
          state_dict_key = ""
          if isinstance(checkpoint, dict):
              if (
                  use_ema
                  and checkpoint.get("state_dict_ema", None) is not None
              ):
                  state_dict_key = "state_dict_ema"
              elif (
                  use_ema
                  and checkpoint.get("model_ema", None) is not None
              ):
                  state_dict_key = "model_ema"
              elif "state_dict" in checkpoint:
                  state_dict_key = "state_dict"
              elif "model" in checkpoint:
                  state_dict_key = "model"
          if state_dict_key:
              state_dict = checkpoint[state_dict_key]
              new_state_dict = OrderedDict()
              for k, v in state_dict.items():
                  # strip `module.` prefix
                  name = k[7:] if k.startswith("module") else k
                  new_state_dict[name] = v
              state_dict = new_state_dict
          else:
              state_dict = checkpoint
          _logger.info(
              "Loaded {} from checkpoint '{}'".format(
                  state_dict_key, checkpoint_path
              )
          )
          return state_dict
      else:
          _logger.error(
              "No checkpoint found at '{}'".format(checkpoint_path)
          )
          raise FileNotFoundError()
  
imports: import logging
  import os
  import math
  from collections import OrderedDict
  from copy import deepcopy
  from typing import Any, Callable, Optional, Tuple
  import torch
  import torch.nn
  from torch.hub import load_state_dict_from_url
  from timm.models.features import FeatureListNet, FeatureDictNet, FeatureHookNet
  from timm.models.fx_features import FeatureGraphNet
  from timm.models.hub import has_hf_hub, download_cached_file, load_state_dict_from_hf
  from timm.models.layers import Conv2dSame, Linear
  from timm.models.vision_transformer import VisionTransformer, resize_pos_embed
  import tutel.impls.moe_layer
label_desc: Use image node_id 1 for calling a global function with example usage: load_state_dict(checkpoint_path, use_ema) and returns: state_dict