code: def state_dict_data_parallel_fix(load_state_dict, curr_state_dict):
      load_keys = list(load_state_dict.keys())
      curr_keys = list(curr_state_dict.keys())
  
      redo_dp = False
      undo_dp = False
      if not curr_keys[0].startswith("module.") and load_keys[
          0
      ].startswith("module."):
          undo_dp = True
      elif curr_keys[0].startswith("module.") and not load_keys[
          0
      ].startswith("module."):
          redo_dp = True
  
      if undo_dp:
          new_state_dict = OrderedDict()
          for k, v in load_state_dict.items():
              name = k[7:]  # remove `module.`
              new_state_dict[name] = v
          # load params
      elif redo_dp:
          new_state_dict = OrderedDict()
          for k, v in load_state_dict.items():
              name = "module." + k  # remove `module.`
              new_state_dict[name] = v
      else:
          new_state_dict = load_state_dict
      return new_state_dict
  
imports: from collections import OrderedDict
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: state_dict_data_parallel_fix(load_state_dict, curr_state_dict) and returns: new_state_dict