code: def convert_to_haiku(a, h, device):
      generator = Generator(h).to(device)
      state_dict_g = load_checkpoint(a.checkpoint_file, device)
      generator.load_state_dict(state_dict_g["generator"])
      generator.eval()
      generator.remove_weight_norm()
      hk_map = {}
      for a, b in generator.state_dict().items():
          print(a, b.shape)
          if a.startswith("conv_pre"):
              a = "generator/~/conv1_d"
          elif a.startswith("conv_post"):
              a = "generator/~/conv1_d_1"
          elif a.startswith("ups."):
              ii = a.split(".")[1]
              a = f"generator/~/ups_{ii}"
          elif a.startswith("resblocks."):
              _, x, y, z, _ = a.split(".")
              a = f"generator/~/res_block1_{x}/~/{y}_{z}"
          print(a, b.shape)
          if a not in hk_map:
              hk_map[a] = {}
          if len(b.shape) == 1:
              hk_map[a]["b"] = b.numpy()
          else:
              if "ups" in a:
                  hk_map[a]["w"] = np.rot90(b.numpy(), k=1, axes=(0, 2))
              elif "conv" in a:
                  hk_map[a]["w"] = np.swapaxes(b.numpy(), 0, 2)
              else:
                  hk_map[a]["w"] = b.numpy()
  
      FLAGS.ckpt_dir.mkdir(parents=True, exist_ok=True)
      with open(FLAGS.ckpt_dir / "hk_hifi.pickle", "wb") as f:
          pickle.dump(hk_map, f)
  
file_path: underthesea/underthesea/pipeline/say/viettts_/hifigan/convert_torch_model_to_haiku.py
imports: import argparse
  import json
  import os
  import pickle
  import numpy
  import torch
  from .config import FLAGS
  from .torch_model import Generator
label_desc: Use image node_id 2 for calling a global function with example usage: convert_to_haiku(a, h, device) without return types