code: def map_state_dict(checkpoint, use_ema=False):
      state_dict_key = "state_dict"
      if isinstance(checkpoint, dict):
          if use_ema and "state_dict_ema" in checkpoint:
              state_dict_key = "state_dict_ema"
      if state_dict_key and state_dict_key in checkpoint:
          new_state_dict = OrderedDict()
          for k, v in checkpoint[state_dict_key].items():
              # strip `module.` prefix
              name = k[7:] if k.startswith("module") else k
              new_state_dict[name] = v
          state_dict = new_state_dict
      elif "model_state" in checkpoint:
          state_dict_key = "model_state"
          new_state_dict = OrderedDict()
          for k, v in checkpoint[state_dict_key].items():
              # strip `model.` prefix
              name = k[6:] if k.startswith("model") else k
              new_state_dict[name] = v
          state_dict = new_state_dict
      else:
          state_dict = checkpoint
      log.info("New state_dict keys: %s", state_dict_key)
      return state_dict
  
file_path: towhee/towhee/models/timesformer/timesformer_utils.py
imports: import logging
  import math
  import subprocess
  from pathlib import Path
  from functools import partial
  from collections import OrderedDict
  import torch
  from torch import nn
  from towhee.models.vit.vit_utils import get_configs
label_desc: Use image node_id 1 for calling a global function with example usage: map_state_dict(checkpoint, use_ema) and returns: state_dict