code: def remove_model_prefix(state_dict):
      # HACK: this is a hack to get the model to load properly, get rid of 'model.' prefix
      for key in list(state_dict["state_dict"].keys()):
          if key.startswith("model."):
              new_key = key[len("model.") :]
              state_dict["state_dict"][new_key] = state_dict[
                  "state_dict"
              ].pop(key)
  
      # HACK: something is wrong with the state dict being loaded...
      return state_dict["state_dict"]
  
file_path: flash-attention/training/src/utils/checkpoint.py
imports: import re
  from pathlib import Path
  import torch
  import math
  from einops import rearrange
label_desc: Use image node_id 4 for calling a global function with example usage: remove_model_prefix(state_dict) and returns: state_dict