code: def transform_state_dict_to_dtype(original_state_dict, dtype="bf16"):
      """
      This function will transform a state dict to bf16
      """
      # following block may fail
      sd_copy = copy.deepcopy(original_state_dict)
      for name in original_state_dict:
          if sd_copy[name].is_floating_point():
              if dtype == "bf16":
                  sd_copy[name] = original_state_dict[name].bfloat16()
              if dtype == "fp32":
                  sd_copy[name] = original_state_dict[name].float()
      return sd_copy
  
imports: import yaml
  from pathlib import Path
  import copy
  import torch
label_desc: Use image node_id 2 for calling a global function with example usage: transform_state_dict_to_dtype(original_state_dict, dtype) and returns: sd_copy