convert.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "709"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#old_path = \"/data1/consistency_models_cifar10/ori_ckpt/cd-l1/checkpoint_80\"\n",
    "#old_path = \"/data1/consistency_models_cifar10/ori_ckpt/cd-l2/checkpoint_80\"\n",
    "#old_path = \"/data1/consistency_models_cifar10/ori_ckpt/cd-lpips/checkpoint_80\"\n",
    "#old_path = \"/data1/consistency_models_cifar10/ori_ckpt/cifar10-continuous-cd-l2/checkpoint_40\"\n",
    "#old_path = \"/data1/consistency_models_cifar10/ori_ckpt/cifar10-continuous-cd-l2-stopgrad/checkpoint_40\"\n",
    "#old_path = \"/data1/consistency_models_cifar10/ori_ckpt/cifar10-continuous-cd-lpips-stopgrad/checkpoint_40\"\n",
    "#old_path = \"/data1/consistency_models_cifar10/ori_ckpt/continuous-ct-l2/checkpoint_80\"\n",
    "#old_path = \"/data1/consistency_models_cifar10/ori_ckpt/continuous-ct-lpips/checkpoint_40\"\n",
    "#old_path = \"/data1/consistency_models_cifar10/ori_ckpt/ct-lpips/checkpoint_74\"\n",
    "old_path = \"/data1/consistency_models_cifar10/ori_ckpt/edm_cifar10_ema/edm_cifar10_ema\"\n",
    "\n",
    "\n",
    "new_path = old_path.replace('ori_ckpt', 'convert_ckpt') + '.pth'\n",
    "import blobfile\n",
    "from flax import serialization\n",
    "from collections import OrderedDict\n",
    "import torch\n",
    "with blobfile.BlobFile(old_path, \"rb\") as fp:\n",
    "        checkpoint_contents = fp.read()\n",
    "        # 从checkpoint_contents中加载字典\n",
    "        loaded_ckpt = serialization.msgpack_restore(checkpoint_contents)\n",
    "\n",
    "        ckpt_dict = OrderedDict()\n",
    "        def print_nested_dict(dictionary, parent_key=\"\", separator=\".\"):\n",
    "                for key, value in dictionary.items():\n",
    "                        new_key = f\"{parent_key}{separator}{key}\" if parent_key else key\n",
    "                        if isinstance(value, dict):\n",
    "                                print_nested_dict(value, new_key, separator)\n",
    "                        else:\n",
    "                                #print(f\"{new_key}:{value.shape}\")\n",
    "                                pass\n",
    "                        ckpt_dict[new_key]= value\n",
    "        print_nested_dict(loaded_ckpt['params_ema'])\n",
    "len(ckpt_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from copy import deepcopy\n",
    "#A = torch.load(\"checkpoint1.pth\")\n",
    "A = ckpt_dict\n",
    "B = torch.load(\"ckpt.pth\")\n",
    "C=deepcopy(B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "odict_keys(['AttnBlockpp_0.GroupNorm_0.bias', 'AttnBlockpp_0.GroupNorm_0.scale', 'AttnBlockpp_0.GroupNorm_0', 'AttnBlockpp_0.NIN_0.W', 'AttnBlockpp_0.NIN_0.b', 'AttnBlockpp_0.NIN_0', 'AttnBlockpp_0.NIN_1.W', 'AttnBlockpp_0.NIN_1.b', 'AttnBlockpp_0.NIN_1', 'AttnBlockpp_0.NIN_2.W', 'AttnBlockpp_0.NIN_2.b', 'AttnBlockpp_0.NIN_2', 'AttnBlockpp_0.NIN_3.W', 'AttnBlockpp_0.NIN_3.b', 'AttnBlockpp_0.NIN_3', 'AttnBlockpp_0', 'AttnBlockpp_1.GroupNorm_0.bias', 'AttnBlockpp_1.GroupNorm_0.scale', 'AttnBlockpp_1.GroupNorm_0', 'AttnBlockpp_1.NIN_0.W', 'AttnBlockpp_1.NIN_0.b', 'AttnBlockpp_1.NIN_0', 'AttnBlockpp_1.NIN_1.W', 'AttnBlockpp_1.NIN_1.b', 'AttnBlockpp_1.NIN_1', 'AttnBlockpp_1.NIN_2.W', 'AttnBlockpp_1.NIN_2.b', 'AttnBlockpp_1.NIN_2', 'AttnBlockpp_1.NIN_3.W', 'AttnBlockpp_1.NIN_3.b', 'AttnBlockpp_1.NIN_3', 'AttnBlockpp_1', 'AttnBlockpp_2.GroupNorm_0.bias', 'AttnBlockpp_2.GroupNorm_0.scale', 'AttnBlockpp_2.GroupNorm_0', 'AttnBlockpp_2.NIN_0.W', 'AttnBlockpp_2.NIN_0.b', 'AttnBlockpp_2.NIN_0', 'AttnBlockpp_2.NIN_1.W', 'AttnBlockpp_2.NIN_1.b', 'AttnBlockpp_2.NIN_1', 'AttnBlockpp_2.NIN_2.W', 'AttnBlockpp_2.NIN_2.b', 'AttnBlockpp_2.NIN_2', 'AttnBlockpp_2.NIN_3.W', 'AttnBlockpp_2.NIN_3.b', 'AttnBlockpp_2.NIN_3', 'AttnBlockpp_2', 'AttnBlockpp_3.GroupNorm_0.bias', 'AttnBlockpp_3.GroupNorm_0.scale', 'AttnBlockpp_3.GroupNorm_0', 'AttnBlockpp_3.NIN_0.W', 'AttnBlockpp_3.NIN_0.b', 'AttnBlockpp_3.NIN_0', 'AttnBlockpp_3.NIN_1.W', 'AttnBlockpp_3.NIN_1.b', 'AttnBlockpp_3.NIN_1', 'AttnBlockpp_3.NIN_2.W', 'AttnBlockpp_3.NIN_2.b', 'AttnBlockpp_3.NIN_2', 'AttnBlockpp_3.NIN_3.W', 'AttnBlockpp_3.NIN_3.b', 'AttnBlockpp_3.NIN_3', 'AttnBlockpp_3', 'AttnBlockpp_4.GroupNorm_0.bias', 'AttnBlockpp_4.GroupNorm_0.scale', 'AttnBlockpp_4.GroupNorm_0', 'AttnBlockpp_4.NIN_0.W', 'AttnBlockpp_4.NIN_0.b', 'AttnBlockpp_4.NIN_0', 'AttnBlockpp_4.NIN_1.W', 'AttnBlockpp_4.NIN_1.b', 'AttnBlockpp_4.NIN_1', 'AttnBlockpp_4.NIN_2.W', 'AttnBlockpp_4.NIN_2.b', 'AttnBlockpp_4.NIN_2', 'AttnBlockpp_4.NIN_3.W', 'AttnBlockpp_4.NIN_3.b', 'AttnBlockpp_4.NIN_3', 'AttnBlockpp_4', 'AttnBlockpp_5.GroupNorm_0.bias', 'AttnBlockpp_5.GroupNorm_0.scale', 'AttnBlockpp_5.GroupNorm_0', 'AttnBlockpp_5.NIN_0.W', 'AttnBlockpp_5.NIN_0.b', 'AttnBlockpp_5.NIN_0', 'AttnBlockpp_5.NIN_1.W', 'AttnBlockpp_5.NIN_1.b', 'AttnBlockpp_5.NIN_1', 'AttnBlockpp_5.NIN_2.W', 'AttnBlockpp_5.NIN_2.b', 'AttnBlockpp_5.NIN_2', 'AttnBlockpp_5.NIN_3.W', 'AttnBlockpp_5.NIN_3.b', 'AttnBlockpp_5.NIN_3', 'AttnBlockpp_5', 'Conv_0.bias', 'Conv_0.kernel', 'Conv_0', 'Conv_1.bias', 'Conv_1.kernel', 'Conv_1', 'Dense_0.bias', 'Dense_0.kernel', 'Dense_0', 'Dense_1.bias', 'Dense_1.kernel', 'Dense_1', 'Downsample_0.Conv2d_0.bias', 'Downsample_0.Conv2d_0.weight', 'Downsample_0.Conv2d_0', 'Downsample_0', 'Downsample_1.Conv2d_0.bias', 'Downsample_1.Conv2d_0.weight', 'Downsample_1.Conv2d_0', 'Downsample_1', 'GaussianFourierProjection_0.W', 'GaussianFourierProjection_0', 'GroupNorm_0.bias', 'GroupNorm_0.scale', 'GroupNorm_0', 'ResnetBlockBigGANpp_0.Conv_0.bias', 'ResnetBlockBigGANpp_0.Conv_0.kernel', 'ResnetBlockBigGANpp_0.Conv_0', 'ResnetBlockBigGANpp_0.Conv_1.bias', 'ResnetBlockBigGANpp_0.Conv_1.kernel', 'ResnetBlockBigGANpp_0.Conv_1', 'ResnetBlockBigGANpp_0.Conv_2.bias', 'ResnetBlockBigGANpp_0.Conv_2.kernel', 'ResnetBlockBigGANpp_0.Conv_2', 'ResnetBlockBigGANpp_0.Dense_0.bias', 'ResnetBlockBigGANpp_0.Dense_0.kernel', 'ResnetBlockBigGANpp_0.Dense_0', 'ResnetBlockBigGANpp_0.GroupNorm_0.bias', 'ResnetBlockBigGANpp_0.GroupNorm_0.scale', 'ResnetBlockBigGANpp_0.GroupNorm_0', 'ResnetBlockBigGANpp_0.GroupNorm_1.bias', 'ResnetBlockBigGANpp_0.GroupNorm_1.scale', 'ResnetBlockBigGANpp_0.GroupNorm_1', 'ResnetBlockBigGANpp_0', 'ResnetBlockBigGANpp_1.Conv_0.bias', 'ResnetBlockBigGANpp_1.Conv_0.kernel', 'ResnetBlockBigGANpp_1.Conv_0', 'ResnetBlockBigGANpp_1.Conv_1.bias', 'ResnetBlockBigGANpp_1.Conv_1.kernel', 'ResnetBlockBigGANpp_1.Conv_1', 'ResnetBlockBigGANpp_1.Dense_0.bias', 'ResnetBlockBigGANpp_1.Dense_0.kernel', 'ResnetBlockBigGANpp_1.Dense_0', 'ResnetBlockBigGANpp_1.GroupNorm_0.bias', 'ResnetBlockBigGANpp_1.GroupNorm_0.scale', 'ResnetBlockBigGANpp_1.GroupNorm_0', 'ResnetBlockBigGANpp_1.GroupNorm_1.bias', 'ResnetBlockBigGANpp_1.GroupNorm_1.scale', 'ResnetBlockBigGANpp_1.GroupNorm_1', 'ResnetBlockBigGANpp_1', 'ResnetBlockBigGANpp_10.Conv_0.bias', 'ResnetBlockBigGANpp_10.Conv_0.kernel', 'ResnetBlockBigGANpp_10.Conv_0', 'ResnetBlockBigGANpp_10.Conv_1.bias', 'ResnetBlockBigGANpp_10.Conv_1.kernel', 'ResnetBlockBigGANpp_10.Conv_1', 'ResnetBlockBigGANpp_10.Dense_0.bias', 'ResnetBlockBigGANpp_10.Dense_0.kernel', 'ResnetBlockBigGANpp_10.Dense_0', 'ResnetBlockBigGANpp_10.GroupNorm_0.bias', 'ResnetBlockBigGANpp_10.GroupNorm_0.scale', 'ResnetBlockBigGANpp_10.GroupNorm_0', 'ResnetBlockBigGANpp_10.GroupNorm_1.bias', 'ResnetBlockBigGANpp_10.GroupNorm_1.scale', 'ResnetBlockBigGANpp_10.GroupNorm_1', 'ResnetBlockBigGANpp_10', 'ResnetBlockBigGANpp_11.Conv_0.bias', 'ResnetBlockBigGANpp_11.Conv_0.kernel', 'ResnetBlockBigGANpp_11.Conv_0', 'ResnetBlockBigGANpp_11.Conv_1.bias', 'ResnetBlockBigGANpp_11.Conv_1.kernel', 'ResnetBlockBigGANpp_11.Conv_1', 'ResnetBlockBigGANpp_11.Dense_0.bias', 'ResnetBlockBigGANpp_11.Dense_0.kernel', 'ResnetBlockBigGANpp_11.Dense_0', 'ResnetBlockBigGANpp_11.GroupNorm_0.bias', 'ResnetBlockBigGANpp_11.GroupNorm_0.scale', 'ResnetBlockBigGANpp_11.GroupNorm_0', 'ResnetBlockBigGANpp_11.GroupNorm_1.bias', 'ResnetBlockBigGANpp_11.GroupNorm_1.scale', 'ResnetBlockBigGANpp_11.GroupNorm_1', 'ResnetBlockBigGANpp_11', 'ResnetBlockBigGANpp_12.Conv_0.bias', 'ResnetBlockBigGANpp_12.Conv_0.kernel', 'ResnetBlockBigGANpp_12.Conv_0', 'ResnetBlockBigGANpp_12.Conv_1.bias', 'ResnetBlockBigGANpp_12.Conv_1.kernel', 'ResnetBlockBigGANpp_12.Conv_1', 'ResnetBlockBigGANpp_12.Dense_0.bias', 'ResnetBlockBigGANpp_12.Dense_0.kernel', 'ResnetBlockBigGANpp_12.Dense_0', 'ResnetBlockBigGANpp_12.GroupNorm_0.bias', 'ResnetBlockBigGANpp_12.GroupNorm_0.scale', 'ResnetBlockBigGANpp_12.GroupNorm_0', 'ResnetBlockBigGANpp_12.GroupNorm_1.bias', 'ResnetBlockBigGANpp_12.GroupNorm_1.scale', 'ResnetBlockBigGANpp_12.GroupNorm_1', 'ResnetBlockBigGANpp_12', 'ResnetBlockBigGANpp_13.Conv_0.bias', 'ResnetBlockBigGANpp_13.Conv_0.kernel', 'ResnetBlockBigGANpp_13.Conv_0', 'ResnetBlockBigGANpp_13.Conv_1.bias', 'ResnetBlockBigGANpp_13.Conv_1.kernel', 'ResnetBlockBigGANpp_13.Conv_1', 'ResnetBlockBigGANpp_13.Dense_0.bias', 'ResnetBlockBigGANpp_13.Dense_0.kernel', 'ResnetBlockBigGANpp_13.Dense_0', 'ResnetBlockBigGANpp_13.GroupNorm_0.bias', 'ResnetBlockBigGANpp_13.GroupNorm_0.scale', 'ResnetBlockBigGANpp_13.GroupNorm_0', 'ResnetBlockBigGANpp_13.GroupNorm_1.bias', 'ResnetBlockBigGANpp_13.GroupNorm_1.scale', 'ResnetBlockBigGANpp_13.GroupNorm_1', 'ResnetBlockBigGANpp_13', 'ResnetBlockBigGANpp_14.Conv_0.bias', 'ResnetBlockBigGANpp_14.Conv_0.kernel', 'ResnetBlockBigGANpp_14.Conv_0', 'ResnetBlockBigGANpp_14.Conv_1.bias', 'ResnetBlockBigGANpp_14.Conv_1.kernel', 'ResnetBlockBigGANpp_14.Conv_1', 'ResnetBlockBigGANpp_14.Dense_0.bias', 'ResnetBlockBigGANpp_14.Dense_0.kernel', 'ResnetBlockBigGANpp_14.Dense_0', 'ResnetBlockBigGANpp_14.GroupNorm_0.bias', 'ResnetBlockBigGANpp_14.GroupNorm_0.scale', 'ResnetBlockBigGANpp_14.GroupNorm_0', 'ResnetBlockBigGANpp_14.GroupNorm_1.bias', 'ResnetBlockBigGANpp_14.GroupNorm_1.scale', 'ResnetBlockBigGANpp_14.GroupNorm_1', 'ResnetBlockBigGANpp_14', 'ResnetBlockBigGANpp_15.Conv_0.bias', 'ResnetBlockBigGANpp_15.Conv_0.kernel', 'ResnetBlockBigGANpp_15.Conv_0', 'ResnetBlockBigGANpp_15.Conv_1.bias', 'ResnetBlockBigGANpp_15.Conv_1.kernel', 'ResnetBlockBigGANpp_15.Conv_1', 'ResnetBlockBigGANpp_15.Dense_0.bias', 'ResnetBlockBigGANpp_15.Dense_0.kernel', 'ResnetBlockBigGANpp_15.Dense_0', 'ResnetBlockBigGANpp_15.GroupNorm_0.bias', 'ResnetBlockBigGANpp_15.GroupNorm_0.scale', 'ResnetBlockBigGANpp_15.GroupNorm_0', 'ResnetBlockBigGANpp_15.GroupNorm_1.bias', 'ResnetBlockBigGANpp_15.GroupNorm_1.scale', 'ResnetBlockBigGANpp_15.GroupNorm_1', 'ResnetBlockBigGANpp_15', 'ResnetBlockBigGANpp_16.Conv_0.bias', 'ResnetBlockBigGANpp_16.Conv_0.kernel', 'ResnetBlockBigGANpp_16.Conv_0', 'ResnetBlockBigGANpp_16.Conv_1.bias', 'ResnetBlockBigGANpp_16.Conv_1.kernel', 'ResnetBlockBigGANpp_16.Conv_1', 'ResnetBlockBigGANpp_16.Conv_2.bias', 'ResnetBlockBigGANpp_16.Conv_2.kernel', 'ResnetBlockBigGANpp_16.Conv_2', 'ResnetBlockBigGANpp_16.Dense_0.bias', 'ResnetBlockBigGANpp_16.Dense_0.kernel', 'ResnetBlockBigGANpp_16.Dense_0', 'ResnetBlockBigGANpp_16.GroupNorm_0.bias', 'ResnetBlockBigGANpp_16.GroupNorm_0.scale', 'ResnetBlockBigGANpp_16.GroupNorm_0', 'ResnetBlockBigGANpp_16.GroupNorm_1.bias', 'ResnetBlockBigGANpp_16.GroupNorm_1.scale', 'ResnetBlockBigGANpp_16.GroupNorm_1', 'ResnetBlockBigGANpp_16', 'ResnetBlockBigGANpp_17.Conv_0.bias', 'ResnetBlockBigGANpp_17.Conv_0.kernel', 'ResnetBlockBigGANpp_17.Conv_0', 'ResnetBlockBigGANpp_17.Conv_1.bias', 'ResnetBlockBigGANpp_17.Conv_1.kernel', 'ResnetBlockBigGANpp_17.Conv_1', 'ResnetBlockBigGANpp_17.Conv_2.bias', 'ResnetBlockBigGANpp_17.Conv_2.kernel', 'ResnetBlockBigGANpp_17.Conv_2', 'ResnetBlockBigGANpp_17.Dense_0.bias', 'ResnetBlockBigGANpp_17.Dense_0.kernel', 'ResnetBlockBigGANpp_17.Dense_0', 'ResnetBlockBigGANpp_17.GroupNorm_0.bias', 'ResnetBlockBigGANpp_17.GroupNorm_0.scale', 'ResnetBlockBigGANpp_17.GroupNorm_0', 'ResnetBlockBigGANpp_17.GroupNorm_1.bias', 'ResnetBlockBigGANpp_17.GroupNorm_1.scale', 'ResnetBlockBigGANpp_17.GroupNorm_1', 'ResnetBlockBigGANpp_17', 'ResnetBlockBigGANpp_18.Conv_0.bias', 'ResnetBlockBigGANpp_18.Conv_0.kernel', 'ResnetBlockBigGANpp_18.Conv_0', 'ResnetBlockBigGANpp_18.Conv_1.bias', 'ResnetBlockBigGANpp_18.Conv_1.kernel', 'ResnetBlockBigGANpp_18.Conv_1', 'ResnetBlockBigGANpp_18.Conv_2.bias', 'ResnetBlockBigGANpp_18.Conv_2.kernel', 'ResnetBlockBigGANpp_18.Conv_2', 'ResnetBlockBigGANpp_18.Dense_0.bias', 'ResnetBlockBigGANpp_18.Dense_0.kernel', 'ResnetBlockBigGANpp_18.Dense_0', 'ResnetBlockBigGANpp_18.GroupNorm_0.bias', 'ResnetBlockBigGANpp_18.GroupNorm_0.scale', 'ResnetBlockBigGANpp_18.GroupNorm_0', 'ResnetBlockBigGANpp_18.GroupNorm_1.bias', 'ResnetBlockBigGANpp_18.GroupNorm_1.scale', 'ResnetBlockBigGANpp_18.GroupNorm_1', 'ResnetBlockBigGANpp_18', 'ResnetBlockBigGANpp_19.Conv_0.bias', 'ResnetBlockBigGANpp_19.Conv_0.kernel', 'ResnetBlockBigGANpp_19.Conv_0', 'ResnetBlockBigGANpp_19.Conv_1.bias', 'ResnetBlockBigGANpp_19.Conv_1.kernel', 'ResnetBlockBigGANpp_19.Conv_1', 'ResnetBlockBigGANpp_19.Conv_2.bias', 'ResnetBlockBigGANpp_19.Conv_2.kernel', 'ResnetBlockBigGANpp_19.Conv_2', 'ResnetBlockBigGANpp_19.Dense_0.bias', 'ResnetBlockBigGANpp_19.Dense_0.kernel', 'ResnetBlockBigGANpp_19.Dense_0', 'ResnetBlockBigGANpp_19.GroupNorm_0.bias', 'ResnetBlockBigGANpp_19.GroupNorm_0.scale', 'ResnetBlockBigGANpp_19.GroupNorm_0', 'ResnetBlockBigGANpp_19.GroupNorm_1.bias', 'ResnetBlockBigGANpp_19.GroupNorm_1.scale', 'ResnetBlockBigGANpp_19.GroupNorm_1', 'ResnetBlockBigGANpp_19', 'ResnetBlockBigGANpp_2.Conv_0.bias', 'ResnetBlockBigGANpp_2.Conv_0.kernel', 'ResnetBlockBigGANpp_2.Conv_0', 'ResnetBlockBigGANpp_2.Conv_1.bias', 'ResnetBlockBigGANpp_2.Conv_1.kernel', 'ResnetBlockBigGANpp_2.Conv_1', 'ResnetBlockBigGANpp_2.Dense_0.bias', 'ResnetBlockBigGANpp_2.Dense_0.kernel', 'ResnetBlockBigGANpp_2.Dense_0', 'ResnetBlockBigGANpp_2.GroupNorm_0.bias', 'ResnetBlockBigGANpp_2.GroupNorm_0.scale', 'ResnetBlockBigGANpp_2.GroupNorm_0', 'ResnetBlockBigGANpp_2.GroupNorm_1.bias', 'ResnetBlockBigGANpp_2.GroupNorm_1.scale', 'ResnetBlockBigGANpp_2.GroupNorm_1', 'ResnetBlockBigGANpp_2', 'ResnetBlockBigGANpp_20.Conv_0.bias', 'ResnetBlockBigGANpp_20.Conv_0.kernel', 'ResnetBlockBigGANpp_20.Conv_0', 'ResnetBlockBigGANpp_20.Conv_1.bias', 'ResnetBlockBigGANpp_20.Conv_1.kernel', 'ResnetBlockBigGANpp_20.Conv_1', 'ResnetBlockBigGANpp_20.Conv_2.bias', 'ResnetBlockBigGANpp_20.Conv_2.kernel', 'ResnetBlockBigGANpp_20.Conv_2', 'ResnetBlockBigGANpp_20.Dense_0.bias', 'ResnetBlockBigGANpp_20.Dense_0.kernel', 'ResnetBlockBigGANpp_20.Dense_0', 'ResnetBlockBigGANpp_20.GroupNorm_0.bias', 'ResnetBlockBigGANpp_20.GroupNorm_0.scale', 'ResnetBlockBigGANpp_20.GroupNorm_0', 'ResnetBlockBigGANpp_20.GroupNorm_1.bias', 'ResnetBlockBigGANpp_20.GroupNorm_1.scale', 'ResnetBlockBigGANpp_20.GroupNorm_1', 'ResnetBlockBigGANpp_20', 'ResnetBlockBigGANpp_21.Conv_0.bias', 'ResnetBlockBigGANpp_21.Conv_0.kernel', 'ResnetBlockBigGANpp_21.Conv_0', 'ResnetBlockBigGANpp_21.Conv_1.bias', 'ResnetBlockBigGANpp_21.Conv_1.kernel', 'ResnetBlockBigGANpp_21.Conv_1', 'ResnetBlockBigGANpp_21.Conv_2.bias', 'ResnetBlockBigGANpp_21.Conv_2.kernel', 'ResnetBlockBigGANpp_21.Conv_2', 'ResnetBlockBigGANpp_21.Dense_0.bias', 'ResnetBlockBigGANpp_21.Dense_0.kernel', 'ResnetBlockBigGANpp_21.Dense_0', 'ResnetBlockBigGANpp_21.GroupNorm_0.bias', 'ResnetBlockBigGANpp_21.GroupNorm_0.scale', 'ResnetBlockBigGANpp_21.GroupNorm_0', 'ResnetBlockBigGANpp_21.GroupNorm_1.bias', 'ResnetBlockBigGANpp_21.GroupNorm_1.scale', 'ResnetBlockBigGANpp_21.GroupNorm_1', 'ResnetBlockBigGANpp_21', 'ResnetBlockBigGANpp_22.Conv_0.bias', 'ResnetBlockBigGANpp_22.Conv_0.kernel', 'ResnetBlockBigGANpp_22.Conv_0', 'ResnetBlockBigGANpp_22.Conv_1.bias', 'ResnetBlockBigGANpp_22.Conv_1.kernel', 'ResnetBlockBigGANpp_22.Conv_1', 'ResnetBlockBigGANpp_22.Conv_2.bias', 'ResnetBlockBigGANpp_22.Conv_2.kernel', 'ResnetBlockBigGANpp_22.Conv_2', 'ResnetBlockBigGANpp_22.Dense_0.bias', 'ResnetBlockBigGANpp_22.Dense_0.kernel', 'ResnetBlockBigGANpp_22.Dense_0', 'ResnetBlockBigGANpp_22.GroupNorm_0.bias', 'ResnetBlockBigGANpp_22.GroupNorm_0.scale', 'ResnetBlockBigGANpp_22.GroupNorm_0', 'ResnetBlockBigGANpp_22.GroupNorm_1.bias', 'ResnetBlockBigGANpp_22.GroupNorm_1.scale', 'ResnetBlockBigGANpp_22.GroupNorm_1', 'ResnetBlockBigGANpp_22', 'ResnetBlockBigGANpp_23.Conv_0.bias', 'ResnetBlockBigGANpp_23.Conv_0.kernel', 'ResnetBlockBigGANpp_23.Conv_0', 'ResnetBlockBigGANpp_23.Conv_1.bias', 'ResnetBlockBigGANpp_23.Conv_1.kernel', 'ResnetBlockBigGANpp_23.Conv_1', 'ResnetBlockBigGANpp_23.Conv_2.bias', 'ResnetBlockBigGANpp_23.Conv_2.kernel', 'ResnetBlockBigGANpp_23.Conv_2', 'ResnetBlockBigGANpp_23.Dense_0.bias', 'ResnetBlockBigGANpp_23.Dense_0.kernel', 'ResnetBlockBigGANpp_23.Dense_0', 'ResnetBlockBigGANpp_23.GroupNorm_0.bias', 'ResnetBlockBigGANpp_23.GroupNorm_0.scale', 'ResnetBlockBigGANpp_23.GroupNorm_0', 'ResnetBlockBigGANpp_23.GroupNorm_1.bias', 'ResnetBlockBigGANpp_23.GroupNorm_1.scale', 'ResnetBlockBigGANpp_23.GroupNorm_1', 'ResnetBlockBigGANpp_23', 'ResnetBlockBigGANpp_24.Conv_0.bias', 'ResnetBlockBigGANpp_24.Conv_0.kernel', 'ResnetBlockBigGANpp_24.Conv_0', 'ResnetBlockBigGANpp_24.Conv_1.bias', 'ResnetBlockBigGANpp_24.Conv_1.kernel', 'ResnetBlockBigGANpp_24.Conv_1', 'ResnetBlockBigGANpp_24.Conv_2.bias', 'ResnetBlockBigGANpp_24.Conv_2.kernel', 'ResnetBlockBigGANpp_24.Conv_2', 'ResnetBlockBigGANpp_24.Dense_0.bias', 'ResnetBlockBigGANpp_24.Dense_0.kernel', 'ResnetBlockBigGANpp_24.Dense_0', 'ResnetBlockBigGANpp_24.GroupNorm_0.bias', 'ResnetBlockBigGANpp_24.GroupNorm_0.scale', 'ResnetBlockBigGANpp_24.GroupNorm_0', 'ResnetBlockBigGANpp_24.GroupNorm_1.bias', 'ResnetBlockBigGANpp_24.GroupNorm_1.scale', 'ResnetBlockBigGANpp_24.GroupNorm_1', 'ResnetBlockBigGANpp_24', 'ResnetBlockBigGANpp_25.Conv_0.bias', 'ResnetBlockBigGANpp_25.Conv_0.kernel', 'ResnetBlockBigGANpp_25.Conv_0', 'ResnetBlockBigGANpp_25.Conv_1.bias', 'ResnetBlockBigGANpp_25.Conv_1.kernel', 'ResnetBlockBigGANpp_25.Conv_1', 'ResnetBlockBigGANpp_25.Conv_2.bias', 'ResnetBlockBigGANpp_25.Conv_2.kernel', 'ResnetBlockBigGANpp_25.Conv_2', 'ResnetBlockBigGANpp_25.Dense_0.bias', 'ResnetBlockBigGANpp_25.Dense_0.kernel', 'ResnetBlockBigGANpp_25.Dense_0', 'ResnetBlockBigGANpp_25.GroupNorm_0.bias', 'ResnetBlockBigGANpp_25.GroupNorm_0.scale', 'ResnetBlockBigGANpp_25.GroupNorm_0', 'ResnetBlockBigGANpp_25.GroupNorm_1.bias', 'ResnetBlockBigGANpp_25.GroupNorm_1.scale', 'ResnetBlockBigGANpp_25.GroupNorm_1', 'ResnetBlockBigGANpp_25', 'ResnetBlockBigGANpp_26.Conv_0.bias', 'ResnetBlockBigGANpp_26.Conv_0.kernel', 'ResnetBlockBigGANpp_26.Conv_0', 'ResnetBlockBigGANpp_26.Conv_1.bias', 'ResnetBlockBigGANpp_26.Conv_1.kernel', 'ResnetBlockBigGANpp_26.Conv_1', 'ResnetBlockBigGANpp_26.Conv_2.bias', 'ResnetBlockBigGANpp_26.Conv_2.kernel', 'ResnetBlockBigGANpp_26.Conv_2', 'ResnetBlockBigGANpp_26.Dense_0.bias', 'ResnetBlockBigGANpp_26.Dense_0.kernel', 'ResnetBlockBigGANpp_26.Dense_0', 'ResnetBlockBigGANpp_26.GroupNorm_0.bias', 'ResnetBlockBigGANpp_26.GroupNorm_0.scale', 'ResnetBlockBigGANpp_26.GroupNorm_0', 'ResnetBlockBigGANpp_26.GroupNorm_1.bias', 'ResnetBlockBigGANpp_26.GroupNorm_1.scale', 'ResnetBlockBigGANpp_26.GroupNorm_1', 'ResnetBlockBigGANpp_26', 'ResnetBlockBigGANpp_27.Conv_0.bias', 'ResnetBlockBigGANpp_27.Conv_0.kernel', 'ResnetBlockBigGANpp_27.Conv_0', 'ResnetBlockBigGANpp_27.Conv_1.bias', 'ResnetBlockBigGANpp_27.Conv_1.kernel', 'ResnetBlockBigGANpp_27.Conv_1', 'ResnetBlockBigGANpp_27.Conv_2.bias', 'ResnetBlockBigGANpp_27.Conv_2.kernel', 'ResnetBlockBigGANpp_27.Conv_2', 'ResnetBlockBigGANpp_27.Dense_0.bias', 'ResnetBlockBigGANpp_27.Dense_0.kernel', 'ResnetBlockBigGANpp_27.Dense_0', 'ResnetBlockBigGANpp_27.GroupNorm_0.bias', 'ResnetBlockBigGANpp_27.GroupNorm_0.scale', 'ResnetBlockBigGANpp_27.GroupNorm_0', 'ResnetBlockBigGANpp_27.GroupNorm_1.bias', 'ResnetBlockBigGANpp_27.GroupNorm_1.scale', 'ResnetBlockBigGANpp_27.GroupNorm_1', 'ResnetBlockBigGANpp_27', 'ResnetBlockBigGANpp_28.Conv_0.bias', 'ResnetBlockBigGANpp_28.Conv_0.kernel', 'ResnetBlockBigGANpp_28.Conv_0', 'ResnetBlockBigGANpp_28.Conv_1.bias', 'ResnetBlockBigGANpp_28.Conv_1.kernel', 'ResnetBlockBigGANpp_28.Conv_1', 'ResnetBlockBigGANpp_28.Conv_2.bias', 'ResnetBlockBigGANpp_28.Conv_2.kernel', 'ResnetBlockBigGANpp_28.Conv_2', 'ResnetBlockBigGANpp_28.Dense_0.bias', 'ResnetBlockBigGANpp_28.Dense_0.kernel', 'ResnetBlockBigGANpp_28.Dense_0', 'ResnetBlockBigGANpp_28.GroupNorm_0.bias', 'ResnetBlockBigGANpp_28.GroupNorm_0.scale', 'ResnetBlockBigGANpp_28.GroupNorm_0', 'ResnetBlockBigGANpp_28.GroupNorm_1.bias', 'ResnetBlockBigGANpp_28.GroupNorm_1.scale', 'ResnetBlockBigGANpp_28.GroupNorm_1', 'ResnetBlockBigGANpp_28', 'ResnetBlockBigGANpp_29.Conv_0.bias', 'ResnetBlockBigGANpp_29.Conv_0.kernel', 'ResnetBlockBigGANpp_29.Conv_0', 'ResnetBlockBigGANpp_29.Conv_1.bias', 'ResnetBlockBigGANpp_29.Conv_1.kernel', 'ResnetBlockBigGANpp_29.Conv_1', 'ResnetBlockBigGANpp_29.Conv_2.bias', 'ResnetBlockBigGANpp_29.Conv_2.kernel', 'ResnetBlockBigGANpp_29.Conv_2', 'ResnetBlockBigGANpp_29.Dense_0.bias', 'ResnetBlockBigGANpp_29.Dense_0.kernel', 'ResnetBlockBigGANpp_29.Dense_0', 'ResnetBlockBigGANpp_29.GroupNorm_0.bias', 'ResnetBlockBigGANpp_29.GroupNorm_0.scale', 'ResnetBlockBigGANpp_29.GroupNorm_0', 'ResnetBlockBigGANpp_29.GroupNorm_1.bias', 'ResnetBlockBigGANpp_29.GroupNorm_1.scale', 'ResnetBlockBigGANpp_29.GroupNorm_1', 'ResnetBlockBigGANpp_29', 'ResnetBlockBigGANpp_3.Conv_0.bias', 'ResnetBlockBigGANpp_3.Conv_0.kernel', 'ResnetBlockBigGANpp_3.Conv_0', 'ResnetBlockBigGANpp_3.Conv_1.bias', 'ResnetBlockBigGANpp_3.Conv_1.kernel', 'ResnetBlockBigGANpp_3.Conv_1', 'ResnetBlockBigGANpp_3.Dense_0.bias', 'ResnetBlockBigGANpp_3.Dense_0.kernel', 'ResnetBlockBigGANpp_3.Dense_0', 'ResnetBlockBigGANpp_3.GroupNorm_0.bias', 'ResnetBlockBigGANpp_3.GroupNorm_0.scale', 'ResnetBlockBigGANpp_3.GroupNorm_0', 'ResnetBlockBigGANpp_3.GroupNorm_1.bias', 'ResnetBlockBigGANpp_3.GroupNorm_1.scale', 'ResnetBlockBigGANpp_3.GroupNorm_1', 'ResnetBlockBigGANpp_3', 'ResnetBlockBigGANpp_30.Conv_0.bias', 'ResnetBlockBigGANpp_30.Conv_0.kernel', 'ResnetBlockBigGANpp_30.Conv_0', 'ResnetBlockBigGANpp_30.Conv_1.bias', 'ResnetBlockBigGANpp_30.Conv_1.kernel', 'ResnetBlockBigGANpp_30.Conv_1', 'ResnetBlockBigGANpp_30.Conv_2.bias', 'ResnetBlockBigGANpp_30.Conv_2.kernel', 'ResnetBlockBigGANpp_30.Conv_2', 'ResnetBlockBigGANpp_30.Dense_0.bias', 'ResnetBlockBigGANpp_30.Dense_0.kernel', 'ResnetBlockBigGANpp_30.Dense_0', 'ResnetBlockBigGANpp_30.GroupNorm_0.bias', 'ResnetBlockBigGANpp_30.GroupNorm_0.scale', 'ResnetBlockBigGANpp_30.GroupNorm_0', 'ResnetBlockBigGANpp_30.GroupNorm_1.bias', 'ResnetBlockBigGANpp_30.GroupNorm_1.scale', 'ResnetBlockBigGANpp_30.GroupNorm_1', 'ResnetBlockBigGANpp_30', 'ResnetBlockBigGANpp_31.Conv_0.bias', 'ResnetBlockBigGANpp_31.Conv_0.kernel', 'ResnetBlockBigGANpp_31.Conv_0', 'ResnetBlockBigGANpp_31.Conv_1.bias', 'ResnetBlockBigGANpp_31.Conv_1.kernel', 'ResnetBlockBigGANpp_31.Conv_1', 'ResnetBlockBigGANpp_31.Conv_2.bias', 'ResnetBlockBigGANpp_31.Conv_2.kernel', 'ResnetBlockBigGANpp_31.Conv_2', 'ResnetBlockBigGANpp_31.Dense_0.bias', 'ResnetBlockBigGANpp_31.Dense_0.kernel', 'ResnetBlockBigGANpp_31.Dense_0', 'ResnetBlockBigGANpp_31.GroupNorm_0.bias', 'ResnetBlockBigGANpp_31.GroupNorm_0.scale', 'ResnetBlockBigGANpp_31.GroupNorm_0', 'ResnetBlockBigGANpp_31.GroupNorm_1.bias', 'ResnetBlockBigGANpp_31.GroupNorm_1.scale', 'ResnetBlockBigGANpp_31.GroupNorm_1', 'ResnetBlockBigGANpp_31', 'ResnetBlockBigGANpp_32.Conv_0.bias', 'ResnetBlockBigGANpp_32.Conv_0.kernel', 'ResnetBlockBigGANpp_32.Conv_0', 'ResnetBlockBigGANpp_32.Conv_1.bias', 'ResnetBlockBigGANpp_32.Conv_1.kernel', 'ResnetBlockBigGANpp_32.Conv_1', 'ResnetBlockBigGANpp_32.Conv_2.bias', 'ResnetBlockBigGANpp_32.Conv_2.kernel', 'ResnetBlockBigGANpp_32.Conv_2', 'ResnetBlockBigGANpp_32.Dense_0.bias', 'ResnetBlockBigGANpp_32.Dense_0.kernel', 'ResnetBlockBigGANpp_32.Dense_0', 'ResnetBlockBigGANpp_32.GroupNorm_0.bias', 'ResnetBlockBigGANpp_32.GroupNorm_0.scale', 'ResnetBlockBigGANpp_32.GroupNorm_0', 'ResnetBlockBigGANpp_32.GroupNorm_1.bias', 'ResnetBlockBigGANpp_32.GroupNorm_1.scale', 'ResnetBlockBigGANpp_32.GroupNorm_1', 'ResnetBlockBigGANpp_32', 'ResnetBlockBigGANpp_4.Conv_0.bias', 'ResnetBlockBigGANpp_4.Conv_0.kernel', 'ResnetBlockBigGANpp_4.Conv_0', 'ResnetBlockBigGANpp_4.Conv_1.bias', 'ResnetBlockBigGANpp_4.Conv_1.kernel', 'ResnetBlockBigGANpp_4.Conv_1', 'ResnetBlockBigGANpp_4.Conv_2.bias', 'ResnetBlockBigGANpp_4.Conv_2.kernel', 'ResnetBlockBigGANpp_4.Conv_2', 'ResnetBlockBigGANpp_4.Dense_0.bias', 'ResnetBlockBigGANpp_4.Dense_0.kernel', 'ResnetBlockBigGANpp_4.Dense_0', 'ResnetBlockBigGANpp_4.GroupNorm_0.bias', 'ResnetBlockBigGANpp_4.GroupNorm_0.scale', 'ResnetBlockBigGANpp_4.GroupNorm_0', 'ResnetBlockBigGANpp_4.GroupNorm_1.bias', 'ResnetBlockBigGANpp_4.GroupNorm_1.scale', 'ResnetBlockBigGANpp_4.GroupNorm_1', 'ResnetBlockBigGANpp_4', 'ResnetBlockBigGANpp_5.Conv_0.bias', 'ResnetBlockBigGANpp_5.Conv_0.kernel', 'ResnetBlockBigGANpp_5.Conv_0', 'ResnetBlockBigGANpp_5.Conv_1.bias', 'ResnetBlockBigGANpp_5.Conv_1.kernel', 'ResnetBlockBigGANpp_5.Conv_1', 'ResnetBlockBigGANpp_5.Dense_0.bias', 'ResnetBlockBigGANpp_5.Dense_0.kernel', 'ResnetBlockBigGANpp_5.Dense_0', 'ResnetBlockBigGANpp_5.GroupNorm_0.bias', 'ResnetBlockBigGANpp_5.GroupNorm_0.scale', 'ResnetBlockBigGANpp_5.GroupNorm_0', 'ResnetBlockBigGANpp_5.GroupNorm_1.bias', 'ResnetBlockBigGANpp_5.GroupNorm_1.scale', 'ResnetBlockBigGANpp_5.GroupNorm_1', 'ResnetBlockBigGANpp_5', 'ResnetBlockBigGANpp_6.Conv_0.bias', 'ResnetBlockBigGANpp_6.Conv_0.kernel', 'ResnetBlockBigGANpp_6.Conv_0', 'ResnetBlockBigGANpp_6.Conv_1.bias', 'ResnetBlockBigGANpp_6.Conv_1.kernel', 'ResnetBlockBigGANpp_6.Conv_1', 'ResnetBlockBigGANpp_6.Dense_0.bias', 'ResnetBlockBigGANpp_6.Dense_0.kernel', 'ResnetBlockBigGANpp_6.Dense_0', 'ResnetBlockBigGANpp_6.GroupNorm_0.bias', 'ResnetBlockBigGANpp_6.GroupNorm_0.scale', 'ResnetBlockBigGANpp_6.GroupNorm_0', 'ResnetBlockBigGANpp_6.GroupNorm_1.bias', 'ResnetBlockBigGANpp_6.GroupNorm_1.scale', 'ResnetBlockBigGANpp_6.GroupNorm_1', 'ResnetBlockBigGANpp_6', 'ResnetBlockBigGANpp_7.Conv_0.bias', 'ResnetBlockBigGANpp_7.Conv_0.kernel', 'ResnetBlockBigGANpp_7.Conv_0', 'ResnetBlockBigGANpp_7.Conv_1.bias', 'ResnetBlockBigGANpp_7.Conv_1.kernel', 'ResnetBlockBigGANpp_7.Conv_1', 'ResnetBlockBigGANpp_7.Dense_0.bias', 'ResnetBlockBigGANpp_7.Dense_0.kernel', 'ResnetBlockBigGANpp_7.Dense_0', 'ResnetBlockBigGANpp_7.GroupNorm_0.bias', 'ResnetBlockBigGANpp_7.GroupNorm_0.scale', 'ResnetBlockBigGANpp_7.GroupNorm_0', 'ResnetBlockBigGANpp_7.GroupNorm_1.bias', 'ResnetBlockBigGANpp_7.GroupNorm_1.scale', 'ResnetBlockBigGANpp_7.GroupNorm_1', 'ResnetBlockBigGANpp_7', 'ResnetBlockBigGANpp_8.Conv_0.bias', 'ResnetBlockBigGANpp_8.Conv_0.kernel', 'ResnetBlockBigGANpp_8.Conv_0', 'ResnetBlockBigGANpp_8.Conv_1.bias', 'ResnetBlockBigGANpp_8.Conv_1.kernel', 'ResnetBlockBigGANpp_8.Conv_1', 'ResnetBlockBigGANpp_8.Dense_0.bias', 'ResnetBlockBigGANpp_8.Dense_0.kernel', 'ResnetBlockBigGANpp_8.Dense_0', 'ResnetBlockBigGANpp_8.GroupNorm_0.bias', 'ResnetBlockBigGANpp_8.GroupNorm_0.scale', 'ResnetBlockBigGANpp_8.GroupNorm_0', 'ResnetBlockBigGANpp_8.GroupNorm_1.bias', 'ResnetBlockBigGANpp_8.GroupNorm_1.scale', 'ResnetBlockBigGANpp_8.GroupNorm_1', 'ResnetBlockBigGANpp_8', 'ResnetBlockBigGANpp_9.Conv_0.bias', 'ResnetBlockBigGANpp_9.Conv_0.kernel', 'ResnetBlockBigGANpp_9.Conv_0', 'ResnetBlockBigGANpp_9.Conv_1.bias', 'ResnetBlockBigGANpp_9.Conv_1.kernel', 'ResnetBlockBigGANpp_9.Conv_1', 'ResnetBlockBigGANpp_9.Conv_2.bias', 'ResnetBlockBigGANpp_9.Conv_2.kernel', 'ResnetBlockBigGANpp_9.Conv_2', 'ResnetBlockBigGANpp_9.Dense_0.bias', 'ResnetBlockBigGANpp_9.Dense_0.kernel', 'ResnetBlockBigGANpp_9.Dense_0', 'ResnetBlockBigGANpp_9.GroupNorm_0.bias', 'ResnetBlockBigGANpp_9.GroupNorm_0.scale', 'ResnetBlockBigGANpp_9.GroupNorm_0', 'ResnetBlockBigGANpp_9.GroupNorm_1.bias', 'ResnetBlockBigGANpp_9.GroupNorm_1.scale', 'ResnetBlockBigGANpp_9.GroupNorm_1', 'ResnetBlockBigGANpp_9'])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "odict_keys(['sigmas', 'all_modules.0.W', 'all_modules.1.weight', 'all_modules.1.bias', 'all_modules.2.weight', 'all_modules.2.bias', 'all_modules.3.weight', 'all_modules.3.bias', 'all_modules.4.GroupNorm_0.weight', 'all_modules.4.GroupNorm_0.bias', 'all_modules.4.Conv_0.weight', 'all_modules.4.Conv_0.bias', 'all_modules.4.Dense_0.weight', 'all_modules.4.Dense_0.bias', 'all_modules.4.GroupNorm_1.weight', 'all_modules.4.GroupNorm_1.bias', 'all_modules.4.Conv_1.weight', 'all_modules.4.Conv_1.bias', 'all_modules.4.Conv_2.weight', 'all_modules.4.Conv_2.bias', 'all_modules.5.GroupNorm_0.weight', 'all_modules.5.GroupNorm_0.bias', 'all_modules.5.Conv_0.weight', 'all_modules.5.Conv_0.bias', 'all_modules.5.Dense_0.weight', 'all_modules.5.Dense_0.bias', 'all_modules.5.GroupNorm_1.weight', 'all_modules.5.GroupNorm_1.bias', 'all_modules.5.Conv_1.weight', 'all_modules.5.Conv_1.bias', 'all_modules.6.GroupNorm_0.weight', 'all_modules.6.GroupNorm_0.bias', 'all_modules.6.Conv_0.weight', 'all_modules.6.Conv_0.bias', 'all_modules.6.Dense_0.weight', 'all_modules.6.Dense_0.bias', 'all_modules.6.GroupNorm_1.weight', 'all_modules.6.GroupNorm_1.bias', 'all_modules.6.Conv_1.weight', 'all_modules.6.Conv_1.bias', 'all_modules.7.GroupNorm_0.weight', 'all_modules.7.GroupNorm_0.bias', 'all_modules.7.Conv_0.weight', 'all_modules.7.Conv_0.bias', 'all_modules.7.Dense_0.weight', 'all_modules.7.Dense_0.bias', 'all_modules.7.GroupNorm_1.weight', 'all_modules.7.GroupNorm_1.bias', 'all_modules.7.Conv_1.weight', 'all_modules.7.Conv_1.bias', 'all_modules.8.GroupNorm_0.weight', 'all_modules.8.GroupNorm_0.bias', 'all_modules.8.Conv_0.weight', 'all_modules.8.Conv_0.bias', 'all_modules.8.Dense_0.weight', 'all_modules.8.Dense_0.bias', 'all_modules.8.GroupNorm_1.weight', 'all_modules.8.GroupNorm_1.bias', 'all_modules.8.Conv_1.weight', 'all_modules.8.Conv_1.bias', 'all_modules.8.Conv_2.weight', 'all_modules.8.Conv_2.bias', 'all_modules.9.Conv2d_0.weight', 'all_modules.9.Conv2d_0.bias', 'all_modules.10.GroupNorm_0.weight', 'all_modules.10.GroupNorm_0.bias', 'all_modules.10.Conv_0.weight', 'all_modules.10.Conv_0.bias', 'all_modules.10.Dense_0.weight', 'all_modules.10.Dense_0.bias', 'all_modules.10.GroupNorm_1.weight', 'all_modules.10.GroupNorm_1.bias', 'all_modules.10.Conv_1.weight', 'all_modules.10.Conv_1.bias', 'all_modules.11.GroupNorm_0.weight', 'all_modules.11.GroupNorm_0.bias', 'all_modules.11.NIN_0.W', 'all_modules.11.NIN_0.b', 'all_modules.11.NIN_1.W', 'all_modules.11.NIN_1.b', 'all_modules.11.NIN_2.W', 'all_modules.11.NIN_2.b', 'all_modules.11.NIN_3.W', 'all_modules.11.NIN_3.b', 'all_modules.12.GroupNorm_0.weight', 'all_modules.12.GroupNorm_0.bias', 'all_modules.12.Conv_0.weight', 'all_modules.12.Conv_0.bias', 'all_modules.12.Dense_0.weight', 'all_modules.12.Dense_0.bias', 'all_modules.12.GroupNorm_1.weight', 'all_modules.12.GroupNorm_1.bias', 'all_modules.12.Conv_1.weight', 'all_modules.12.Conv_1.bias', 'all_modules.13.GroupNorm_0.weight', 'all_modules.13.GroupNorm_0.bias', 'all_modules.13.NIN_0.W', 'all_modules.13.NIN_0.b', 'all_modules.13.NIN_1.W', 'all_modules.13.NIN_1.b', 'all_modules.13.NIN_2.W', 'all_modules.13.NIN_2.b', 'all_modules.13.NIN_3.W', 'all_modules.13.NIN_3.b', 'all_modules.14.GroupNorm_0.weight', 'all_modules.14.GroupNorm_0.bias', 'all_modules.14.Conv_0.weight', 'all_modules.14.Conv_0.bias', 'all_modules.14.Dense_0.weight', 'all_modules.14.Dense_0.bias', 'all_modules.14.GroupNorm_1.weight', 'all_modules.14.GroupNorm_1.bias', 'all_modules.14.Conv_1.weight', 'all_modules.14.Conv_1.bias', 'all_modules.15.GroupNorm_0.weight', 'all_modules.15.GroupNorm_0.bias', 'all_modules.15.NIN_0.W', 'all_modules.15.NIN_0.b', 'all_modules.15.NIN_1.W', 'all_modules.15.NIN_1.b', 'all_modules.15.NIN_2.W', 'all_modules.15.NIN_2.b', 'all_modules.15.NIN_3.W', 'all_modules.15.NIN_3.b', 'all_modules.16.GroupNorm_0.weight', 'all_modules.16.GroupNorm_0.bias', 'all_modules.16.Conv_0.weight', 'all_modules.16.Conv_0.bias', 'all_modules.16.Dense_0.weight', 'all_modules.16.Dense_0.bias', 'all_modules.16.GroupNorm_1.weight', 'all_modules.16.GroupNorm_1.bias', 'all_modules.16.Conv_1.weight', 'all_modules.16.Conv_1.bias', 'all_modules.17.GroupNorm_0.weight', 'all_modules.17.GroupNorm_0.bias', 'all_modules.17.NIN_0.W', 'all_modules.17.NIN_0.b', 'all_modules.17.NIN_1.W', 'all_modules.17.NIN_1.b', 'all_modules.17.NIN_2.W', 'all_modules.17.NIN_2.b', 'all_modules.17.NIN_3.W', 'all_modules.17.NIN_3.b', 'all_modules.18.GroupNorm_0.weight', 'all_modules.18.GroupNorm_0.bias', 'all_modules.18.Conv_0.weight', 'all_modules.18.Conv_0.bias', 'all_modules.18.Dense_0.weight', 'all_modules.18.Dense_0.bias', 'all_modules.18.GroupNorm_1.weight', 'all_modules.18.GroupNorm_1.bias', 'all_modules.18.Conv_1.weight', 'all_modules.18.Conv_1.bias', 'all_modules.18.Conv_2.weight', 'all_modules.18.Conv_2.bias', 'all_modules.19.Conv2d_0.weight', 'all_modules.19.Conv2d_0.bias', 'all_modules.20.GroupNorm_0.weight', 'all_modules.20.GroupNorm_0.bias', 'all_modules.20.Conv_0.weight', 'all_modules.20.Conv_0.bias', 'all_modules.20.Dense_0.weight', 'all_modules.20.Dense_0.bias', 'all_modules.20.GroupNorm_1.weight', 'all_modules.20.GroupNorm_1.bias', 'all_modules.20.Conv_1.weight', 'all_modules.20.Conv_1.bias', 'all_modules.21.GroupNorm_0.weight', 'all_modules.21.GroupNorm_0.bias', 'all_modules.21.Conv_0.weight', 'all_modules.21.Conv_0.bias', 'all_modules.21.Dense_0.weight', 'all_modules.21.Dense_0.bias', 'all_modules.21.GroupNorm_1.weight', 'all_modules.21.GroupNorm_1.bias', 'all_modules.21.Conv_1.weight', 'all_modules.21.Conv_1.bias', 'all_modules.22.GroupNorm_0.weight', 'all_modules.22.GroupNorm_0.bias', 'all_modules.22.Conv_0.weight', 'all_modules.22.Conv_0.bias', 'all_modules.22.Dense_0.weight', 'all_modules.22.Dense_0.bias', 'all_modules.22.GroupNorm_1.weight', 'all_modules.22.GroupNorm_1.bias', 'all_modules.22.Conv_1.weight', 'all_modules.22.Conv_1.bias', 'all_modules.23.GroupNorm_0.weight', 'all_modules.23.GroupNorm_0.bias', 'all_modules.23.Conv_0.weight', 'all_modules.23.Conv_0.bias', 'all_modules.23.Dense_0.weight', 'all_modules.23.Dense_0.bias', 'all_modules.23.GroupNorm_1.weight', 'all_modules.23.GroupNorm_1.bias', 'all_modules.23.Conv_1.weight', 'all_modules.23.Conv_1.bias', 'all_modules.24.GroupNorm_0.weight', 'all_modules.24.GroupNorm_0.bias', 'all_modules.24.Conv_0.weight', 'all_modules.24.Conv_0.bias', 'all_modules.24.Dense_0.weight', 'all_modules.24.Dense_0.bias', 'all_modules.24.GroupNorm_1.weight', 'all_modules.24.GroupNorm_1.bias', 'all_modules.24.Conv_1.weight', 'all_modules.24.Conv_1.bias', 'all_modules.25.GroupNorm_0.weight', 'all_modules.25.GroupNorm_0.bias', 'all_modules.25.NIN_0.W', 'all_modules.25.NIN_0.b', 'all_modules.25.NIN_1.W', 'all_modules.25.NIN_1.b', 'all_modules.25.NIN_2.W', 'all_modules.25.NIN_2.b', 'all_modules.25.NIN_3.W', 'all_modules.25.NIN_3.b', 'all_modules.26.GroupNorm_0.weight', 'all_modules.26.GroupNorm_0.bias', 'all_modules.26.Conv_0.weight', 'all_modules.26.Conv_0.bias', 'all_modules.26.Dense_0.weight', 'all_modules.26.Dense_0.bias', 'all_modules.26.GroupNorm_1.weight', 'all_modules.26.GroupNorm_1.bias', 'all_modules.26.Conv_1.weight', 'all_modules.26.Conv_1.bias', 'all_modules.27.GroupNorm_0.weight', 'all_modules.27.GroupNorm_0.bias', 'all_modules.27.Conv_0.weight', 'all_modules.27.Conv_0.bias', 'all_modules.27.Dense_0.weight', 'all_modules.27.Dense_0.bias', 'all_modules.27.GroupNorm_1.weight', 'all_modules.27.GroupNorm_1.bias', 'all_modules.27.Conv_1.weight', 'all_modules.27.Conv_1.bias', 'all_modules.27.Conv_2.weight', 'all_modules.27.Conv_2.bias', 'all_modules.28.GroupNorm_0.weight', 'all_modules.28.GroupNorm_0.bias', 'all_modules.28.Conv_0.weight', 'all_modules.28.Conv_0.bias', 'all_modules.28.Dense_0.weight', 'all_modules.28.Dense_0.bias', 'all_modules.28.GroupNorm_1.weight', 'all_modules.28.GroupNorm_1.bias', 'all_modules.28.Conv_1.weight', 'all_modules.28.Conv_1.bias', 'all_modules.28.Conv_2.weight', 'all_modules.28.Conv_2.bias', 'all_modules.29.GroupNorm_0.weight', 'all_modules.29.GroupNorm_0.bias', 'all_modules.29.Conv_0.weight', 'all_modules.29.Conv_0.bias', 'all_modules.29.Dense_0.weight', 'all_modules.29.Dense_0.bias', 'all_modules.29.GroupNorm_1.weight', 'all_modules.29.GroupNorm_1.bias', 'all_modules.29.Conv_1.weight', 'all_modules.29.Conv_1.bias', 'all_modules.29.Conv_2.weight', 'all_modules.29.Conv_2.bias', 'all_modules.30.GroupNorm_0.weight', 'all_modules.30.GroupNorm_0.bias', 'all_modules.30.Conv_0.weight', 'all_modules.30.Conv_0.bias', 'all_modules.30.Dense_0.weight', 'all_modules.30.Dense_0.bias', 'all_modules.30.GroupNorm_1.weight', 'all_modules.30.GroupNorm_1.bias', 'all_modules.30.Conv_1.weight', 'all_modules.30.Conv_1.bias', 'all_modules.30.Conv_2.weight', 'all_modules.30.Conv_2.bias', 'all_modules.31.GroupNorm_0.weight', 'all_modules.31.GroupNorm_0.bias', 'all_modules.31.Conv_0.weight', 'all_modules.31.Conv_0.bias', 'all_modules.31.Dense_0.weight', 'all_modules.31.Dense_0.bias', 'all_modules.31.GroupNorm_1.weight', 'all_modules.31.GroupNorm_1.bias', 'all_modules.31.Conv_1.weight', 'all_modules.31.Conv_1.bias', 'all_modules.31.Conv_2.weight', 'all_modules.31.Conv_2.bias', 'all_modules.32.GroupNorm_0.weight', 'all_modules.32.GroupNorm_0.bias', 'all_modules.32.Conv_0.weight', 'all_modules.32.Conv_0.bias', 'all_modules.32.Dense_0.weight', 'all_modules.32.Dense_0.bias', 'all_modules.32.GroupNorm_1.weight', 'all_modules.32.GroupNorm_1.bias', 'all_modules.32.Conv_1.weight', 'all_modules.32.Conv_1.bias', 'all_modules.32.Conv_2.weight', 'all_modules.32.Conv_2.bias', 'all_modules.33.GroupNorm_0.weight', 'all_modules.33.GroupNorm_0.bias', 'all_modules.33.Conv_0.weight', 'all_modules.33.Conv_0.bias', 'all_modules.33.Dense_0.weight', 'all_modules.33.Dense_0.bias', 'all_modules.33.GroupNorm_1.weight', 'all_modules.33.GroupNorm_1.bias', 'all_modules.33.Conv_1.weight', 'all_modules.33.Conv_1.bias', 'all_modules.33.Conv_2.weight', 'all_modules.33.Conv_2.bias', 'all_modules.34.GroupNorm_0.weight', 'all_modules.34.GroupNorm_0.bias', 'all_modules.34.Conv_0.weight', 'all_modules.34.Conv_0.bias', 'all_modules.34.Dense_0.weight', 'all_modules.34.Dense_0.bias', 'all_modules.34.GroupNorm_1.weight', 'all_modules.34.GroupNorm_1.bias', 'all_modules.34.Conv_1.weight', 'all_modules.34.Conv_1.bias', 'all_modules.34.Conv_2.weight', 'all_modules.34.Conv_2.bias', 'all_modules.35.GroupNorm_0.weight', 'all_modules.35.GroupNorm_0.bias', 'all_modules.35.Conv_0.weight', 'all_modules.35.Conv_0.bias', 'all_modules.35.Dense_0.weight', 'all_modules.35.Dense_0.bias', 'all_modules.35.GroupNorm_1.weight', 'all_modules.35.GroupNorm_1.bias', 'all_modules.35.Conv_1.weight', 'all_modules.35.Conv_1.bias', 'all_modules.35.Conv_2.weight', 'all_modules.35.Conv_2.bias', 'all_modules.36.GroupNorm_0.weight', 'all_modules.36.GroupNorm_0.bias', 'all_modules.36.Conv_0.weight', 'all_modules.36.Conv_0.bias', 'all_modules.36.Dense_0.weight', 'all_modules.36.Dense_0.bias', 'all_modules.36.GroupNorm_1.weight', 'all_modules.36.GroupNorm_1.bias', 'all_modules.36.Conv_1.weight', 'all_modules.36.Conv_1.bias', 'all_modules.36.Conv_2.weight', 'all_modules.36.Conv_2.bias', 'all_modules.37.GroupNorm_0.weight', 'all_modules.37.GroupNorm_0.bias', 'all_modules.37.Conv_0.weight', 'all_modules.37.Conv_0.bias', 'all_modules.37.Dense_0.weight', 'all_modules.37.Dense_0.bias', 'all_modules.37.GroupNorm_1.weight', 'all_modules.37.GroupNorm_1.bias', 'all_modules.37.Conv_1.weight', 'all_modules.37.Conv_1.bias', 'all_modules.37.Conv_2.weight', 'all_modules.37.Conv_2.bias', 'all_modules.38.GroupNorm_0.weight', 'all_modules.38.GroupNorm_0.bias', 'all_modules.38.NIN_0.W', 'all_modules.38.NIN_0.b', 'all_modules.38.NIN_1.W', 'all_modules.38.NIN_1.b', 'all_modules.38.NIN_2.W', 'all_modules.38.NIN_2.b', 'all_modules.38.NIN_3.W', 'all_modules.38.NIN_3.b', 'all_modules.39.GroupNorm_0.weight', 'all_modules.39.GroupNorm_0.bias', 'all_modules.39.Conv_0.weight', 'all_modules.39.Conv_0.bias', 'all_modules.39.Dense_0.weight', 'all_modules.39.Dense_0.bias', 'all_modules.39.GroupNorm_1.weight', 'all_modules.39.GroupNorm_1.bias', 'all_modules.39.Conv_1.weight', 'all_modules.39.Conv_1.bias', 'all_modules.39.Conv_2.weight', 'all_modules.39.Conv_2.bias', 'all_modules.40.GroupNorm_0.weight', 'all_modules.40.GroupNorm_0.bias', 'all_modules.40.Conv_0.weight', 'all_modules.40.Conv_0.bias', 'all_modules.40.Dense_0.weight', 'all_modules.40.Dense_0.bias', 'all_modules.40.GroupNorm_1.weight', 'all_modules.40.GroupNorm_1.bias', 'all_modules.40.Conv_1.weight', 'all_modules.40.Conv_1.bias', 'all_modules.40.Conv_2.weight', 'all_modules.40.Conv_2.bias', 'all_modules.41.GroupNorm_0.weight', 'all_modules.41.GroupNorm_0.bias', 'all_modules.41.Conv_0.weight', 'all_modules.41.Conv_0.bias', 'all_modules.41.Dense_0.weight', 'all_modules.41.Dense_0.bias', 'all_modules.41.GroupNorm_1.weight', 'all_modules.41.GroupNorm_1.bias', 'all_modules.41.Conv_1.weight', 'all_modules.41.Conv_1.bias', 'all_modules.41.Conv_2.weight', 'all_modules.41.Conv_2.bias', 'all_modules.42.GroupNorm_0.weight', 'all_modules.42.GroupNorm_0.bias', 'all_modules.42.Conv_0.weight', 'all_modules.42.Conv_0.bias', 'all_modules.42.Dense_0.weight', 'all_modules.42.Dense_0.bias', 'all_modules.42.GroupNorm_1.weight', 'all_modules.42.GroupNorm_1.bias', 'all_modules.42.Conv_1.weight', 'all_modules.42.Conv_1.bias', 'all_modules.42.Conv_2.weight', 'all_modules.42.Conv_2.bias', 'all_modules.43.GroupNorm_0.weight', 'all_modules.43.GroupNorm_0.bias', 'all_modules.43.Conv_0.weight', 'all_modules.43.Conv_0.bias', 'all_modules.43.Dense_0.weight', 'all_modules.43.Dense_0.bias', 'all_modules.43.GroupNorm_1.weight', 'all_modules.43.GroupNorm_1.bias', 'all_modules.43.Conv_1.weight', 'all_modules.43.Conv_1.bias', 'all_modules.43.Conv_2.weight', 'all_modules.43.Conv_2.bias', 'all_modules.44.GroupNorm_0.weight', 'all_modules.44.GroupNorm_0.bias', 'all_modules.44.Conv_0.weight', 'all_modules.44.Conv_0.bias', 'all_modules.44.Dense_0.weight', 'all_modules.44.Dense_0.bias', 'all_modules.44.GroupNorm_1.weight', 'all_modules.44.GroupNorm_1.bias', 'all_modules.44.Conv_1.weight', 'all_modules.44.Conv_1.bias', 'all_modules.44.Conv_2.weight', 'all_modules.44.Conv_2.bias', 'all_modules.45.weight', 'all_modules.45.bias', 'all_modules.46.weight', 'all_modules.46.bias'])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "odict_keys(['sigmas', 'all_modules.0.W', 'all_modules.1.weight', 'all_modules.1.bias', 'all_modules.2.weight', 'all_modules.2.bias', 'all_modules.3.weight', 'all_modules.3.bias', 'all_modules.4.GroupNorm_0.weight', 'all_modules.4.GroupNorm_0.bias', 'all_modules.4.Conv_0.weight', 'all_modules.4.Conv_0.bias', 'all_modules.4.Dense_0.weight', 'all_modules.4.Dense_0.bias', 'all_modules.4.GroupNorm_1.weight', 'all_modules.4.GroupNorm_1.bias', 'all_modules.4.Conv_1.weight', 'all_modules.4.Conv_1.bias', 'all_modules.4.Conv_2.weight', 'all_modules.4.Conv_2.bias', 'all_modules.5.GroupNorm_0.weight', 'all_modules.5.GroupNorm_0.bias', 'all_modules.5.Conv_0.weight', 'all_modules.5.Conv_0.bias', 'all_modules.5.Dense_0.weight', 'all_modules.5.Dense_0.bias', 'all_modules.5.GroupNorm_1.weight', 'all_modules.5.GroupNorm_1.bias', 'all_modules.5.Conv_1.weight', 'all_modules.5.Conv_1.bias', 'all_modules.6.GroupNorm_0.weight', 'all_modules.6.GroupNorm_0.bias', 'all_modules.6.Conv_0.weight', 'all_modules.6.Conv_0.bias', 'all_modules.6.Dense_0.weight', 'all_modules.6.Dense_0.bias', 'all_modules.6.GroupNorm_1.weight', 'all_modules.6.GroupNorm_1.bias', 'all_modules.6.Conv_1.weight', 'all_modules.6.Conv_1.bias', 'all_modules.7.GroupNorm_0.weight', 'all_modules.7.GroupNorm_0.bias', 'all_modules.7.Conv_0.weight', 'all_modules.7.Conv_0.bias', 'all_modules.7.Dense_0.weight', 'all_modules.7.Dense_0.bias', 'all_modules.7.GroupNorm_1.weight', 'all_modules.7.GroupNorm_1.bias', 'all_modules.7.Conv_1.weight', 'all_modules.7.Conv_1.bias', 'all_modules.8.GroupNorm_0.weight', 'all_modules.8.GroupNorm_0.bias', 'all_modules.8.Conv_0.weight', 'all_modules.8.Conv_0.bias', 'all_modules.8.Dense_0.weight', 'all_modules.8.Dense_0.bias', 'all_modules.8.GroupNorm_1.weight', 'all_modules.8.GroupNorm_1.bias', 'all_modules.8.Conv_1.weight', 'all_modules.8.Conv_1.bias', 'all_modules.8.Conv_2.weight', 'all_modules.8.Conv_2.bias', 'all_modules.9.Conv2d_0.weight', 'all_modules.9.Conv2d_0.bias', 'all_modules.10.GroupNorm_0.weight', 'all_modules.10.GroupNorm_0.bias', 'all_modules.10.Conv_0.weight', 'all_modules.10.Conv_0.bias', 'all_modules.10.Dense_0.weight', 'all_modules.10.Dense_0.bias', 'all_modules.10.GroupNorm_1.weight', 'all_modules.10.GroupNorm_1.bias', 'all_modules.10.Conv_1.weight', 'all_modules.10.Conv_1.bias', 'all_modules.11.GroupNorm_0.weight', 'all_modules.11.GroupNorm_0.bias', 'all_modules.11.NIN_0.W', 'all_modules.11.NIN_0.b', 'all_modules.11.NIN_1.W', 'all_modules.11.NIN_1.b', 'all_modules.11.NIN_2.W', 'all_modules.11.NIN_2.b', 'all_modules.11.NIN_3.W', 'all_modules.11.NIN_3.b', 'all_modules.12.GroupNorm_0.weight', 'all_modules.12.GroupNorm_0.bias', 'all_modules.12.Conv_0.weight', 'all_modules.12.Conv_0.bias', 'all_modules.12.Dense_0.weight', 'all_modules.12.Dense_0.bias', 'all_modules.12.GroupNorm_1.weight', 'all_modules.12.GroupNorm_1.bias', 'all_modules.12.Conv_1.weight', 'all_modules.12.Conv_1.bias', 'all_modules.13.GroupNorm_0.weight', 'all_modules.13.GroupNorm_0.bias', 'all_modules.13.NIN_0.W', 'all_modules.13.NIN_0.b', 'all_modules.13.NIN_1.W', 'all_modules.13.NIN_1.b', 'all_modules.13.NIN_2.W', 'all_modules.13.NIN_2.b', 'all_modules.13.NIN_3.W', 'all_modules.13.NIN_3.b', 'all_modules.14.GroupNorm_0.weight', 'all_modules.14.GroupNorm_0.bias', 'all_modules.14.Conv_0.weight', 'all_modules.14.Conv_0.bias', 'all_modules.14.Dense_0.weight', 'all_modules.14.Dense_0.bias', 'all_modules.14.GroupNorm_1.weight', 'all_modules.14.GroupNorm_1.bias', 'all_modules.14.Conv_1.weight', 'all_modules.14.Conv_1.bias', 'all_modules.15.GroupNorm_0.weight', 'all_modules.15.GroupNorm_0.bias', 'all_modules.15.NIN_0.W', 'all_modules.15.NIN_0.b', 'all_modules.15.NIN_1.W', 'all_modules.15.NIN_1.b', 'all_modules.15.NIN_2.W', 'all_modules.15.NIN_2.b', 'all_modules.15.NIN_3.W', 'all_modules.15.NIN_3.b', 'all_modules.16.GroupNorm_0.weight', 'all_modules.16.GroupNorm_0.bias', 'all_modules.16.Conv_0.weight', 'all_modules.16.Conv_0.bias', 'all_modules.16.Dense_0.weight', 'all_modules.16.Dense_0.bias', 'all_modules.16.GroupNorm_1.weight', 'all_modules.16.GroupNorm_1.bias', 'all_modules.16.Conv_1.weight', 'all_modules.16.Conv_1.bias', 'all_modules.17.GroupNorm_0.weight', 'all_modules.17.GroupNorm_0.bias', 'all_modules.17.NIN_0.W', 'all_modules.17.NIN_0.b', 'all_modules.17.NIN_1.W', 'all_modules.17.NIN_1.b', 'all_modules.17.NIN_2.W', 'all_modules.17.NIN_2.b', 'all_modules.17.NIN_3.W', 'all_modules.17.NIN_3.b', 'all_modules.18.GroupNorm_0.weight', 'all_modules.18.GroupNorm_0.bias', 'all_modules.18.Conv_0.weight', 'all_modules.18.Conv_0.bias', 'all_modules.18.Dense_0.weight', 'all_modules.18.Dense_0.bias', 'all_modules.18.GroupNorm_1.weight', 'all_modules.18.GroupNorm_1.bias', 'all_modules.18.Conv_1.weight', 'all_modules.18.Conv_1.bias', 'all_modules.18.Conv_2.weight', 'all_modules.18.Conv_2.bias', 'all_modules.19.Conv2d_0.weight', 'all_modules.19.Conv2d_0.bias', 'all_modules.20.GroupNorm_0.weight', 'all_modules.20.GroupNorm_0.bias', 'all_modules.20.Conv_0.weight', 'all_modules.20.Conv_0.bias', 'all_modules.20.Dense_0.weight', 'all_modules.20.Dense_0.bias', 'all_modules.20.GroupNorm_1.weight', 'all_modules.20.GroupNorm_1.bias', 'all_modules.20.Conv_1.weight', 'all_modules.20.Conv_1.bias', 'all_modules.21.GroupNorm_0.weight', 'all_modules.21.GroupNorm_0.bias', 'all_modules.21.Conv_0.weight', 'all_modules.21.Conv_0.bias', 'all_modules.21.Dense_0.weight', 'all_modules.21.Dense_0.bias', 'all_modules.21.GroupNorm_1.weight', 'all_modules.21.GroupNorm_1.bias', 'all_modules.21.Conv_1.weight', 'all_modules.21.Conv_1.bias', 'all_modules.22.GroupNorm_0.weight', 'all_modules.22.GroupNorm_0.bias', 'all_modules.22.Conv_0.weight', 'all_modules.22.Conv_0.bias', 'all_modules.22.Dense_0.weight', 'all_modules.22.Dense_0.bias', 'all_modules.22.GroupNorm_1.weight', 'all_modules.22.GroupNorm_1.bias', 'all_modules.22.Conv_1.weight', 'all_modules.22.Conv_1.bias', 'all_modules.23.GroupNorm_0.weight', 'all_modules.23.GroupNorm_0.bias', 'all_modules.23.Conv_0.weight', 'all_modules.23.Conv_0.bias', 'all_modules.23.Dense_0.weight', 'all_modules.23.Dense_0.bias', 'all_modules.23.GroupNorm_1.weight', 'all_modules.23.GroupNorm_1.bias', 'all_modules.23.Conv_1.weight', 'all_modules.23.Conv_1.bias', 'all_modules.24.GroupNorm_0.weight', 'all_modules.24.GroupNorm_0.bias', 'all_modules.24.Conv_0.weight', 'all_modules.24.Conv_0.bias', 'all_modules.24.Dense_0.weight', 'all_modules.24.Dense_0.bias', 'all_modules.24.GroupNorm_1.weight', 'all_modules.24.GroupNorm_1.bias', 'all_modules.24.Conv_1.weight', 'all_modules.24.Conv_1.bias', 'all_modules.25.GroupNorm_0.weight', 'all_modules.25.GroupNorm_0.bias', 'all_modules.25.NIN_0.W', 'all_modules.25.NIN_0.b', 'all_modules.25.NIN_1.W', 'all_modules.25.NIN_1.b', 'all_modules.25.NIN_2.W', 'all_modules.25.NIN_2.b', 'all_modules.25.NIN_3.W', 'all_modules.25.NIN_3.b', 'all_modules.26.GroupNorm_0.weight', 'all_modules.26.GroupNorm_0.bias', 'all_modules.26.Conv_0.weight', 'all_modules.26.Conv_0.bias', 'all_modules.26.Dense_0.weight', 'all_modules.26.Dense_0.bias', 'all_modules.26.GroupNorm_1.weight', 'all_modules.26.GroupNorm_1.bias', 'all_modules.26.Conv_1.weight', 'all_modules.26.Conv_1.bias', 'all_modules.27.GroupNorm_0.weight', 'all_modules.27.GroupNorm_0.bias', 'all_modules.27.Conv_0.weight', 'all_modules.27.Conv_0.bias', 'all_modules.27.Dense_0.weight', 'all_modules.27.Dense_0.bias', 'all_modules.27.GroupNorm_1.weight', 'all_modules.27.GroupNorm_1.bias', 'all_modules.27.Conv_1.weight', 'all_modules.27.Conv_1.bias', 'all_modules.27.Conv_2.weight', 'all_modules.27.Conv_2.bias', 'all_modules.28.GroupNorm_0.weight', 'all_modules.28.GroupNorm_0.bias', 'all_modules.28.Conv_0.weight', 'all_modules.28.Conv_0.bias', 'all_modules.28.Dense_0.weight', 'all_modules.28.Dense_0.bias', 'all_modules.28.GroupNorm_1.weight', 'all_modules.28.GroupNorm_1.bias', 'all_modules.28.Conv_1.weight', 'all_modules.28.Conv_1.bias', 'all_modules.28.Conv_2.weight', 'all_modules.28.Conv_2.bias', 'all_modules.29.GroupNorm_0.weight', 'all_modules.29.GroupNorm_0.bias', 'all_modules.29.Conv_0.weight', 'all_modules.29.Conv_0.bias', 'all_modules.29.Dense_0.weight', 'all_modules.29.Dense_0.bias', 'all_modules.29.GroupNorm_1.weight', 'all_modules.29.GroupNorm_1.bias', 'all_modules.29.Conv_1.weight', 'all_modules.29.Conv_1.bias', 'all_modules.29.Conv_2.weight', 'all_modules.29.Conv_2.bias', 'all_modules.30.GroupNorm_0.weight', 'all_modules.30.GroupNorm_0.bias', 'all_modules.30.Conv_0.weight', 'all_modules.30.Conv_0.bias', 'all_modules.30.Dense_0.weight', 'all_modules.30.Dense_0.bias', 'all_modules.30.GroupNorm_1.weight', 'all_modules.30.GroupNorm_1.bias', 'all_modules.30.Conv_1.weight', 'all_modules.30.Conv_1.bias', 'all_modules.30.Conv_2.weight', 'all_modules.30.Conv_2.bias', 'all_modules.31.GroupNorm_0.weight', 'all_modules.31.GroupNorm_0.bias', 'all_modules.31.Conv_0.weight', 'all_modules.31.Conv_0.bias', 'all_modules.31.Dense_0.weight', 'all_modules.31.Dense_0.bias', 'all_modules.31.GroupNorm_1.weight', 'all_modules.31.GroupNorm_1.bias', 'all_modules.31.Conv_1.weight', 'all_modules.31.Conv_1.bias', 'all_modules.31.Conv_2.weight', 'all_modules.31.Conv_2.bias', 'all_modules.32.GroupNorm_0.weight', 'all_modules.32.GroupNorm_0.bias', 'all_modules.32.Conv_0.weight', 'all_modules.32.Conv_0.bias', 'all_modules.32.Dense_0.weight', 'all_modules.32.Dense_0.bias', 'all_modules.32.GroupNorm_1.weight', 'all_modules.32.GroupNorm_1.bias', 'all_modules.32.Conv_1.weight', 'all_modules.32.Conv_1.bias', 'all_modules.32.Conv_2.weight', 'all_modules.32.Conv_2.bias', 'all_modules.33.GroupNorm_0.weight', 'all_modules.33.GroupNorm_0.bias', 'all_modules.33.Conv_0.weight', 'all_modules.33.Conv_0.bias', 'all_modules.33.Dense_0.weight', 'all_modules.33.Dense_0.bias', 'all_modules.33.GroupNorm_1.weight', 'all_modules.33.GroupNorm_1.bias', 'all_modules.33.Conv_1.weight', 'all_modules.33.Conv_1.bias', 'all_modules.33.Conv_2.weight', 'all_modules.33.Conv_2.bias', 'all_modules.34.GroupNorm_0.weight', 'all_modules.34.GroupNorm_0.bias', 'all_modules.34.Conv_0.weight', 'all_modules.34.Conv_0.bias', 'all_modules.34.Dense_0.weight', 'all_modules.34.Dense_0.bias', 'all_modules.34.GroupNorm_1.weight', 'all_modules.34.GroupNorm_1.bias', 'all_modules.34.Conv_1.weight', 'all_modules.34.Conv_1.bias', 'all_modules.34.Conv_2.weight', 'all_modules.34.Conv_2.bias', 'all_modules.35.GroupNorm_0.weight', 'all_modules.35.GroupNorm_0.bias', 'all_modules.35.Conv_0.weight', 'all_modules.35.Conv_0.bias', 'all_modules.35.Dense_0.weight', 'all_modules.35.Dense_0.bias', 'all_modules.35.GroupNorm_1.weight', 'all_modules.35.GroupNorm_1.bias', 'all_modules.35.Conv_1.weight', 'all_modules.35.Conv_1.bias', 'all_modules.35.Conv_2.weight', 'all_modules.35.Conv_2.bias', 'all_modules.36.GroupNorm_0.weight', 'all_modules.36.GroupNorm_0.bias', 'all_modules.36.Conv_0.weight', 'all_modules.36.Conv_0.bias', 'all_modules.36.Dense_0.weight', 'all_modules.36.Dense_0.bias', 'all_modules.36.GroupNorm_1.weight', 'all_modules.36.GroupNorm_1.bias', 'all_modules.36.Conv_1.weight', 'all_modules.36.Conv_1.bias', 'all_modules.36.Conv_2.weight', 'all_modules.36.Conv_2.bias', 'all_modules.37.GroupNorm_0.weight', 'all_modules.37.GroupNorm_0.bias', 'all_modules.37.Conv_0.weight', 'all_modules.37.Conv_0.bias', 'all_modules.37.Dense_0.weight', 'all_modules.37.Dense_0.bias', 'all_modules.37.GroupNorm_1.weight', 'all_modules.37.GroupNorm_1.bias', 'all_modules.37.Conv_1.weight', 'all_modules.37.Conv_1.bias', 'all_modules.37.Conv_2.weight', 'all_modules.37.Conv_2.bias', 'all_modules.38.GroupNorm_0.weight', 'all_modules.38.GroupNorm_0.bias', 'all_modules.38.NIN_0.W', 'all_modules.38.NIN_0.b', 'all_modules.38.NIN_1.W', 'all_modules.38.NIN_1.b', 'all_modules.38.NIN_2.W', 'all_modules.38.NIN_2.b', 'all_modules.38.NIN_3.W', 'all_modules.38.NIN_3.b', 'all_modules.39.GroupNorm_0.weight', 'all_modules.39.GroupNorm_0.bias', 'all_modules.39.Conv_0.weight', 'all_modules.39.Conv_0.bias', 'all_modules.39.Dense_0.weight', 'all_modules.39.Dense_0.bias', 'all_modules.39.GroupNorm_1.weight', 'all_modules.39.GroupNorm_1.bias', 'all_modules.39.Conv_1.weight', 'all_modules.39.Conv_1.bias', 'all_modules.39.Conv_2.weight', 'all_modules.39.Conv_2.bias', 'all_modules.40.GroupNorm_0.weight', 'all_modules.40.GroupNorm_0.bias', 'all_modules.40.Conv_0.weight', 'all_modules.40.Conv_0.bias', 'all_modules.40.Dense_0.weight', 'all_modules.40.Dense_0.bias', 'all_modules.40.GroupNorm_1.weight', 'all_modules.40.GroupNorm_1.bias', 'all_modules.40.Conv_1.weight', 'all_modules.40.Conv_1.bias', 'all_modules.40.Conv_2.weight', 'all_modules.40.Conv_2.bias', 'all_modules.41.GroupNorm_0.weight', 'all_modules.41.GroupNorm_0.bias', 'all_modules.41.Conv_0.weight', 'all_modules.41.Conv_0.bias', 'all_modules.41.Dense_0.weight', 'all_modules.41.Dense_0.bias', 'all_modules.41.GroupNorm_1.weight', 'all_modules.41.GroupNorm_1.bias', 'all_modules.41.Conv_1.weight', 'all_modules.41.Conv_1.bias', 'all_modules.41.Conv_2.weight', 'all_modules.41.Conv_2.bias', 'all_modules.42.GroupNorm_0.weight', 'all_modules.42.GroupNorm_0.bias', 'all_modules.42.Conv_0.weight', 'all_modules.42.Conv_0.bias', 'all_modules.42.Dense_0.weight', 'all_modules.42.Dense_0.bias', 'all_modules.42.GroupNorm_1.weight', 'all_modules.42.GroupNorm_1.bias', 'all_modules.42.Conv_1.weight', 'all_modules.42.Conv_1.bias', 'all_modules.42.Conv_2.weight', 'all_modules.42.Conv_2.bias', 'all_modules.43.GroupNorm_0.weight', 'all_modules.43.GroupNorm_0.bias', 'all_modules.43.Conv_0.weight', 'all_modules.43.Conv_0.bias', 'all_modules.43.Dense_0.weight', 'all_modules.43.Dense_0.bias', 'all_modules.43.GroupNorm_1.weight', 'all_modules.43.GroupNorm_1.bias', 'all_modules.43.Conv_1.weight', 'all_modules.43.Conv_1.bias', 'all_modules.43.Conv_2.weight', 'all_modules.43.Conv_2.bias', 'all_modules.44.GroupNorm_0.weight', 'all_modules.44.GroupNorm_0.bias', 'all_modules.44.Conv_0.weight', 'all_modules.44.Conv_0.bias', 'all_modules.44.Dense_0.weight', 'all_modules.44.Dense_0.bias', 'all_modules.44.GroupNorm_1.weight', 'all_modules.44.GroupNorm_1.bias', 'all_modules.44.Conv_1.weight', 'all_modules.44.Conv_1.bias', 'all_modules.44.Conv_2.weight', 'all_modules.44.Conv_2.bias', 'all_modules.45.weight', 'all_modules.45.bias', 'all_modules.46.weight', 'all_modules.46.bias'])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Tensor"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(B['all_modules.4.Conv_2.weight'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def back(x):\n",
    "    #print(*reversed(tuple(range(x.dim()))))\n",
    "    return x.permute(*reversed(tuple(range(x.dim()))))\n",
    "def back_conv(x):\n",
    "    #print(*reversed(tuple(range(x.dim()))))\n",
    "    return x.permute(3,2,0,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_10 = dict({'1': 5, '2': 6, '3': 7, '5': 10, '6': 12, '7': 14, '8': 16, '10': 20, '11': 21, '12': 22, '13': 23, '14': 24, '15': 26})\n",
    "c=dict({\"0\":11,\"1\":13,\"2\":15,\"3\":17,\"4\":25,\"5\":38})\n",
    "b_12 = dict({'0': 4, '4': 8, '9': 18, '16': 27, '17': 28, '18': 29, '19': 30, '20': 31, '21': 32, '22': 33, '23': 34, '24': 35, '25': 36, '26': 37, '27': 39, '28': 40, '29': 41, '30': 42, '31': 43, '32': 44})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3270785/2184132048.py:4: UserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means writing to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:206.)\n",
      "  C[f\"all_modules.{y}.GroupNorm_0.weight\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.GroupNorm_0.scale\"]))\n"
     ]
    }
   ],
   "source": [
    "for x,y in b_12.items():\n",
    "    x=int(x)\n",
    "    y=int(y)\n",
    "    C[f\"all_modules.{y}.GroupNorm_0.weight\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.GroupNorm_0.scale\"]))\n",
    "    C[f\"all_modules.{y}.GroupNorm_0.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.GroupNorm_0.bias\"]))\n",
    "    C[f\"all_modules.{y}.Conv_0.weight\"]=back_conv(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Conv_0.kernel\"]))\n",
    "    C[f\"all_modules.{y}.Conv_0.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Conv_0.bias\"]))\n",
    "    C[f\"all_modules.{y}.Dense_0.weight\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Dense_0.kernel\"]))\n",
    "    C[f\"all_modules.{y}.Dense_0.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Dense_0.bias\"]))\n",
    "    C[f\"all_modules.{y}.GroupNorm_1.weight\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.GroupNorm_1.scale\"]))\n",
    "    C[f\"all_modules.{y}.GroupNorm_1.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.GroupNorm_1.bias\"]))\n",
    "    C[f\"all_modules.{y}.Conv_1.weight\"]=back_conv(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Conv_1.kernel\"]))\n",
    "    C[f\"all_modules.{y}.Conv_1.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Conv_1.bias\"]))\n",
    "    C[f\"all_modules.{y}.Conv_2.weight\"]=back_conv(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Conv_2.kernel\"]))\n",
    "    C[f\"all_modules.{y}.Conv_2.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Conv_2.bias\"]))\n",
    "for x,y in a_10.items():\n",
    "    x=int(x)\n",
    "    y=int(y)\n",
    "    C[f\"all_modules.{y}.GroupNorm_0.weight\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.GroupNorm_0.scale\"]))\n",
    "    C[f\"all_modules.{y}.GroupNorm_0.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.GroupNorm_0.bias\"]))\n",
    "    C[f\"all_modules.{y}.Conv_0.weight\"]=back_conv(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Conv_0.kernel\"]))\n",
    "    C[f\"all_modules.{y}.Conv_0.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Conv_0.bias\"]))\n",
    "    C[f\"all_modules.{y}.Dense_0.weight\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Dense_0.kernel\"]))\n",
    "    C[f\"all_modules.{y}.Dense_0.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Dense_0.bias\"]))\n",
    "    C[f\"all_modules.{y}.GroupNorm_1.weight\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.GroupNorm_1.scale\"]))\n",
    "    C[f\"all_modules.{y}.GroupNorm_1.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.GroupNorm_1.bias\"]))\n",
    "    C[f\"all_modules.{y}.Conv_1.weight\"]=back_conv(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Conv_1.kernel\"]))\n",
    "    C[f\"all_modules.{y}.Conv_1.bias\"]=back(torch.from_numpy(A[f\"ResnetBlockBigGANpp_{x}.Conv_1.bias\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "for x,y in c.items():\n",
    "    x=int(x)\n",
    "    y=int(y)\n",
    "    C[f\"all_modules.{y}.GroupNorm_0.weight\"]=back(torch.from_numpy(A[f\"AttnBlockpp_{x}.GroupNorm_0.scale\"]))\n",
    "    C[f\"all_modules.{y}.GroupNorm_0.bias\"]=back(torch.from_numpy(A[f\"AttnBlockpp_{x}.GroupNorm_0.bias\"]))\n",
    "    for cccc in range(4):\n",
    "        C[f\"all_modules.{y}.NIN_{cccc}.W\"]=torch.from_numpy(A[f\"AttnBlockpp_{x}.NIN_{cccc}.W\"])\n",
    "        C[f\"all_modules.{y}.NIN_{cccc}.b\"]=torch.from_numpy(A[f\"AttnBlockpp_{x}.NIN_{cccc}.b\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "C[\"all_modules.0.W\"]=back(torch.from_numpy(A[\"GaussianFourierProjection_0.W\"]))\n",
    "C[\"all_modules.1.weight\"]=back(torch.from_numpy(A[\"Dense_0.kernel\"]))\n",
    "C[\"all_modules.1.bias\"]=back(torch.from_numpy(A[\"Dense_0.bias\"]))\n",
    "C[\"all_modules.2.weight\"]=back(torch.from_numpy(A[\"Dense_1.kernel\"]))\n",
    "C[\"all_modules.2.bias\"]=back(torch.from_numpy(A[\"Dense_1.bias\"]))\n",
    "C[\"all_modules.3.weight\"]=back_conv(torch.from_numpy(A[\"Conv_0.kernel\"]))\n",
    "C[\"all_modules.3.bias\"]=back(torch.from_numpy(A[\"Conv_0.bias\"]))\n",
    "\n",
    "C[\"all_modules.45.weight\"]=back(torch.from_numpy(A[\"GroupNorm_0.scale\"]))\n",
    "C[\"all_modules.45.bias\"]=back(torch.from_numpy(A[\"GroupNorm_0.bias\"]))\n",
    "C[\"all_modules.46.weight\"]=back_conv(torch.from_numpy(A[\"Conv_1.kernel\"]))\n",
    "C[\"all_modules.46.bias\"]=back(torch.from_numpy(A[\"Conv_1.bias\"]))\n",
    "\n",
    "C[\"all_modules.9.Conv2d_0.weight\"]=back_conv(torch.from_numpy(A[\"Downsample_0.Conv2d_0.weight\"]))\n",
    "C[\"all_modules.9.Conv2d_0.bias\"]=back(torch.from_numpy(A[\"Downsample_0.Conv2d_0.bias\"]))\n",
    "C[\"all_modules.19.Conv2d_0.weight\"]=back_conv(torch.from_numpy(A[\"Downsample_1.Conv2d_0.weight\"]))\n",
    "C[\"all_modules.19.Conv2d_0.bias\"]=back(torch.from_numpy(A[\"Downsample_1.Conv2d_0.bias\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-19.2643,   1.5175,   9.8834,  21.5128,   2.0145, -18.8356,   6.0850,\n",
       "          4.8258,  14.9231, -21.0671,   8.1425,   9.8985,  -7.6452,   6.1850,\n",
       "        -21.9804,  -7.8546,  13.1654, -14.4788, -12.0730,  -3.7347,  26.8905,\n",
       "        -33.0785,  22.9105, -10.7170, -23.6326,  -6.1353,   1.1644,  17.0903,\n",
       "        -18.6139, -16.4002,  -0.8027,  -2.1112,  22.0649,  30.6345,  -8.4047,\n",
       "          4.9921, -16.1030,  29.3569,  18.8985,  -3.5282,   9.0032,  13.4016,\n",
       "          3.3607,  17.2367,  15.9929,  -3.1429, -20.4104,  16.4431,  15.6330,\n",
       "         22.6655,   7.7387, -32.8621,  16.6384,   1.9347, -21.7612,  -5.3607,\n",
       "          4.4274,  11.7961,  17.5778,  16.5485,  24.7009,  19.1643, -25.8576,\n",
       "         14.6646, -23.0816,  15.2318,  10.0603,   5.5396, -18.7672, -15.6804,\n",
       "         -4.3610,  13.1263,  38.9641,  12.1402, -19.4894,  23.2745, -18.7513,\n",
       "          3.4976, -16.4826,  18.9395,   6.4940,  -1.3373,  26.7928,   8.1603,\n",
       "        -25.8580,  -6.6758,   0.1610, -43.0085,   9.7588,  -9.6417,   2.2592,\n",
       "        -12.8034,  -1.4683,  11.4133,   5.8974,  29.0003,  -2.3967, -30.8309,\n",
       "        -12.8086,  13.1907,   0.6203,  25.3083, -11.5564,   6.9026,  -9.4813,\n",
       "        -17.9300,  -1.0351,  -8.0281,  15.5218,   0.1055,   8.8926, -20.0540,\n",
       "         -1.8329, -14.1139,   0.7993, -10.2413,  -9.7675, -17.3919,   5.0734,\n",
       "          7.2308,  17.4513, -49.1078,  -9.9800,   6.8105, -11.9240, -29.5836,\n",
       "          3.2388,   5.1486])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C[\"all_modules.0.W\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.makedirs(os.path.dirname(new_path), exist_ok=True)\n",
    "torch.save(C,new_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pt2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}