diff --git "a/diffusion.ipynb" "b/diffusion.ipynb"
--- "a/diffusion.ipynb"
+++ "b/diffusion.ipynb"
@@ -67,6 +67,30 @@
     "from huggingface_hub import notebook_login"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b7a51a73994a43178f1baa379210c892",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "notebook_login()"
+   ]
+  },
   {
    "attachments": {},
    "cell_type": "markdown",
@@ -89,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -196,7 +220,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -227,7 +251,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -244,10 +268,10 @@
     "    # repeat = 2\n",
     "\n",
     "    # dim = 2\n",
-    "    dim = 2\n",
-    "    stride = (2,2) if dim == 2 else (2,2,4)\n",
+    "    dim = 3\n",
+    "    stride = (2,2) if dim == 2 else (2,2,2)\n",
     "    num_image = 2000#32000#20000#15000#7000#25600#3000#10000#1000#10000#5000#2560#800#2560\n",
-    "    batch_size = 50#20#2#100 # 10\n",
+    "    batch_size = 2#50#20#2#100 # 10\n",
     "    n_epoch = 50#20#20#2#5#25 # 120\n",
     "    HII_DIM = 64\n",
     "    num_redshift = 64#512#256#256#64#512#128\n",
@@ -304,7 +328,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -529,37 +553,13 @@
     "# print(\"device =\", config.device)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c74ef5393df14c0b9f283a0dfad7a9e8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "notebook_login()"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
-    "num_image_list = [1600,3200,6400,12800,25600]"
+    "num_image_list = [4]#[1600,3200,6400,12800,25600]"
    ]
   },
   {
@@ -571,16 +571,15 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Number of parameters for nn_model: 111048705\n",
-      "---------------- num_image = 1600 ----------------\n",
-      "run_name = 0702-1718\n",
+      "Number of parameters for nn_model: 306285057\n",
+      "----------------- num_image = 4 ------------------\n",
+      "run_name = 0705-1053\n",
       "Launching training on one GPU.\n",
       "dataset content: <KeysViewHDF5 ['brightness_temp', 'density', 'kwargs', 'params', 'redshifts_distances', 'seeds', 'xH_box']>\n",
       "51200 images can be loaded\n",
       "field.shape = (64, 64, 514)\n",
       "params keys = [b'ION_Tvir_MIN', b'HII_EFF_FACTOR']\n",
-      "loading 1600 images randomly\n",
-      "images loaded: (1600, 1, 64, 64)\n"
+      "loading 4 images randomly\n"
      ]
     },
     {
@@ -594,20 +593,21 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "params loaded: (1600, 2)\n",
-      "images rescaled to [-1.0, 1.119088888168335]\n",
-      "params rescaled to [0.0001848356100438764, 0.9999958400767995]\n"
+      "images loaded: (4, 1, 64, 64, 64)\n",
+      "params loaded: (4, 2)\n",
+      "images rescaled to [-1.0, 1.049072027206421]\n",
+      "params rescaled to [0.02179058530500466, 0.9278468466439764]\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fe5b9721d5d743e4a7628ad9232bf72d",
+       "model_id": "2ecbe26474ef446498cae235911aeb71",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -616,12 +616,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "561f62b97ede4f4dbd5153cfed56b72d",
+       "model_id": "c7e4b83cd28e418ba485f3cbb621bbfc",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -630,12 +630,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2df65647d9dd4843990231ebc184344a",
+       "model_id": "5131eb83f4a5481f86dd5e760994777f",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -644,12 +644,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "98b05c506a6c4e2e9ba5c9cbafd11818",
+       "model_id": "911f1534f52b491e8a1d32333ca197fd",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -658,12 +658,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a2273b8655d74b3e8785f518a098f63e",
+       "model_id": "86c1521c0527401bb3109596384a2c0a",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -672,12 +672,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5107917d669a42279e050372954b845e",
+       "model_id": "3d398aaa68ee4ca9a49edd104836a9ae",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -686,12 +686,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b4f77a0574f04d3991adbd62e962f80f",
+       "model_id": "9862503132bf477b828e6d8dfd34ddff",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -700,12 +700,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "68c69cac2d6f40c4af53f8a31caac370",
+       "model_id": "d6bfe318f0404914934a3f29b77af475",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -714,12 +714,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "374801feaad14d0c824bff08f00c34d6",
+       "model_id": "31af611de29d4cea9403997ed099ec6b",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -728,12 +728,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e9a2f5c8b8614161b96397806d24408f",
+       "model_id": "39d76b3a6ece4e20b7b12d3a1f1bd7cb",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -742,12 +742,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5efd0a2aed1748be824f7c0ec194bbf1",
+       "model_id": "aa1ffafe98544080991f9b16a69b6237",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -756,12 +756,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "11763f061b8c4f85a2e8c01169e7ba29",
+       "model_id": "de8d6d0ca17f44aaadd1cac178fd1382",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -770,12 +770,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0a51162af4754dfeb56ad6576951f6f3",
+       "model_id": "4180238b5ca4458b8d32bb84ebafd6bc",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -784,12 +784,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "011cd27257aa488d9ed6418cae7252d7",
+       "model_id": "c119c446ecce4c48929b84c4289eb902",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -798,12 +798,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "268d518ad6b24bbcb91c0ba4cad79ce5",
+       "model_id": "efa57bccdae94728beb48c25944343df",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -812,12 +812,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ffede80b20774f8d95ed105a171829c7",
+       "model_id": "69122a78260c41998a99114d6b16cb05",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -826,12 +826,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0908d553e2dc474ba1e17f37b3ec0ca2",
+       "model_id": "2dec8b90a16e4fcf97e57086075fc5f4",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -840,12 +840,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f0839a380799488398c37161f1c33c32",
+       "model_id": "553466be83994a14bb115fa4ef492fb9",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -854,12 +854,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "641575d247dc4f19821713d857836bbc",
+       "model_id": "9b2cc8d228d74034804d2dcc0d4cf0c7",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -868,12 +868,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "14815c4292be41b6922bba0423cac1c1",
+       "model_id": "f62e293b34da434e8642429d7f8ee9f1",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -882,12 +882,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9c480c8e167d407da6b860940c9f4df0",
+       "model_id": "cb978aee9b994953b049315b067ca689",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -896,12 +896,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c335431face94892a55146524933a2a8",
+       "model_id": "a9facb2eb98b4f1cb330658ff017d6e2",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -910,12 +910,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6c65c569a9ef4900ba0fe20c1141d203",
+       "model_id": "1e9e7b0ec426476b86e6ad8e0d7f3f6c",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -924,12 +924,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c168f618534b4924b6b2ee11596c5d19",
+       "model_id": "993b95c4019d4c3fba6a921244cd1f85",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -938,12 +938,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b358120c2fb34869aad9840b5e52a48a",
+       "model_id": "102b96fe42a549d491db1b1ef06b2b71",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -952,12 +952,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0d45b8d356a74764afb0a09a4f7e7b38",
+       "model_id": "bea9b3b7ce674edfab82e722729b5b89",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -966,12 +966,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cb5c02e1bbf34b2fa5bb29dda3f483c8",
+       "model_id": "9bd60ffe82bd4905bacc2b679731405d",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -980,12 +980,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b59572b300484b58816a213e8ebf7ecb",
+       "model_id": "bdaff07deceb41609994255b73c32cd9",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -994,12 +994,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "73c59564808e490e9e1a6ff166365f41",
+       "model_id": "192999ef9c8f47c9ad406842de8d4e8d",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1008,12 +1008,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "97a8d0968f65411b9c52718ab344a680",
+       "model_id": "611701080c1f47e1888f853cb1ceb126",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1022,12 +1022,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5e0fbab35c2a48b9add70cc6690960f5",
+       "model_id": "4b00e996c78f4389a6331543159d2e1a",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1036,12 +1036,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e39a6332264344a19450ea9304ce7d8c",
+       "model_id": "42701fb451c64cb593f4fca4d143751c",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1050,12 +1050,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6b969529bd24430c96371972808c8e67",
+       "model_id": "47669f94caa546a78ab26957098d2092",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1064,12 +1064,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cd736ff127f547d589b773892bd3ff08",
+       "model_id": "8afc2c87a8ff4ba3af430245a087f67a",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1078,12 +1078,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a4ca3c9c4e18413bac155a69f6bf967a",
+       "model_id": "7e1ea4c40e9142a9b8559e491aff2529",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1092,12 +1092,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "70adcf6937104ba3b80a6a41953318fe",
+       "model_id": "77015601a26d4fd2a3a41d6d80538d37",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1106,12 +1106,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7c49dd288f6e4d8a85e5e9b15f4dd3d2",
+       "model_id": "b1a35ee7663e49fc9b2386059b53e8c8",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1120,12 +1120,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c105d25a15de4c8a9de4fe83ce671f57",
+       "model_id": "9bbdc5fe2e584026856fb9a2ce6eb1b7",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1134,12 +1134,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "40223545bf004d6288f354696bbd9bee",
+       "model_id": "1cb8d4301ed446b1ad56554ff080560a",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1148,12 +1148,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e696b9312dd44739af9bdd218e1ffa0a",
+       "model_id": "e0c4bee9568447b78db716554a3342b2",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1162,12 +1162,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9f9ff4a333fa48b3bb500b8c6dd0e546",
+       "model_id": "bdd4db66e7344f2f9ba41d1512e76458",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1176,12 +1176,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "80bee9dff0554e2690910c54fd099604",
+       "model_id": "55ceda2803084fc0bb46524e0d586602",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1190,12 +1190,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "595f8b30b2c54010bc45f906e5fd7148",
+       "model_id": "ae7215706afa4ddfa3a674dd76ee8350",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1204,12 +1204,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "98172d80b3a84c63836139c2cafe47d1",
+       "model_id": "512249484a9d45b69dea4ed238f3da30",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1218,12 +1218,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "04b845a13bb6422d9bad982ff365006e",
+       "model_id": "4ec68a068a604ef5b11c99993b8baace",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1232,12 +1232,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "517580f0ffbb4fe5b2650de5355471f1",
+       "model_id": "f1d060aae57a47edbb55b9c19413a77c",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1246,12 +1246,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fac1c16c12bf4b86a5719314acefb248",
+       "model_id": "f848c117255e48dea811b18a159a7175",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1260,12 +1260,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "95533b6eafd14d4bbd133c36e71ea99e",
+       "model_id": "3b0c9aa8cca34000a81abf133d710ccf",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1274,12 +1274,12 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b53949253982428fb092be9bb55a8b89",
+       "model_id": "474eaaaa3ba340b0af823d796ddf7ea0",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
@@ -1288,3050 +1288,503 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f31938c7618c43dbadcfe7b382e92c46",
+       "model_id": "c546e85d0c0a46da9a50001bc63c2f92",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/32 [00:00<?, ?it/s]"
+       "  0%|          | 0/2 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
-    },
+    }
+   ],
+   "source": [
+    "if __name__ == \"__main__\":\n",
+    "    # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)\n",
+    "    config = TrainConfig()\n",
+    "    for i, num_image in enumerate(num_image_list):\n",
+    "        config.num_image = num_image\n",
+    "        ddpm21cm = DDPM21CM(config)\n",
+    "        print(f\" num_image = {ddpm21cm.config.num_image} \".center(50, '-'))\n",
+    "        print(f\"run_name = {ddpm21cm.config.run_name}\")\n",
+    "        notebook_launcher(ddpm21cm.train, num_processes=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "saved model at ./outputs/model_state-N1600\n",
-      "Number of parameters for nn_model: 111048705\n",
-      "---------------- num_image = 3200 ----------------\n",
-      "run_name = 0702-1725\n",
-      "Launching training on one GPU.\n",
-      "dataset content: <KeysViewHDF5 ['brightness_temp', 'density', 'kwargs', 'params', 'redshifts_distances', 'seeds', 'xH_box']>\n",
-      "51200 images can be loaded\n",
-      "field.shape = (64, 64, 514)\n",
-      "params keys = [b'ION_Tvir_MIN', b'HII_EFF_FACTOR']\n",
-      "loading 3200 images randomly\n",
-      "images loaded: (3200, 1, 64, 64)\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Detected kernel version 3.10.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
+      "total 4.4G\n",
+      "-rw-r--r--  1 bxia34  13M Jul  2 21:45 Tvir4.800000190734863-zeta131.34100341796875-N1600.npy\n",
+      "-rw-r--r--  1 bxia34  13M Jul  2 21:26 Tvir5.4770002365112305-zeta200.0-N1600.npy\n",
+      "-rw-r--r--  1 bxia34  13M Jul  2 21:08 Tvir4.698999881744385-zeta30.0-N1600.npy\n",
+      "-rw-r--r--  1 bxia34  13M Jul  2 20:49 Tvir5.599999904632568-zeta19.03700065612793-N1600.npy\n",
+      "-rw-r--r--  1 bxia34  13M Jul  2 20:31 Tvir4.400000095367432-zeta131.34100341796875-N1600.npy\n",
+      "-rw-r--r--  1 bxia34 848M Jul  2 20:13 model_state-N25600\n",
+      "drwxr-xr-x 15 bxia34 4.0K Jul  2 19:09 \u001b[0m\u001b[01;34mlogs\u001b[0m/\n",
+      "-rw-r--r--  1 bxia34 848M Jul  2 18:45 model_state-N12800\n",
+      "-rw-r--r--  1 bxia34 848M Jul  2 18:01 model_state-N6400\n",
+      "-rw-r--r--  1 bxia34 848M Jul  2 17:37 model_state-N3200\n",
+      "-rw-r--r--  1 bxia34 848M Jul  2 17:25 model_state-N1600\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 12:31 Tvir4.800000190734863-zeta131.34100341796875-N2000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 12:12 Tvir5.4770002365112305-zeta200.0-N2000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 11:54 Tvir4.698999881744385-zeta30.0-N2000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 11:35 Tvir5.599999904632568-zeta19.03700065612793-N2000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 11:17 Tvir4.400000095367432-zeta131.34100341796875-N2000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 10:37 Tvir4.800000190734863-zeta131.34100341796875-N32000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 04:47 Tvir5.4770002365112305-zeta200.0-N32000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 04:29 Tvir4.698999881744385-zeta30.0-N32000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 04:11 Tvir5.599999904632568-zeta19.03700065612793-N32000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 03:53 Tvir4.400000095367432-zeta131.34100341796875-N32000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 00:23 Tvir4.800000190734863-zeta131.34100341796875-N20000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jul  1 00:05 Tvir5.4770002365112305-zeta200.0-N20000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 30 23:47 Tvir4.698999881744385-zeta30.0-N20000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 30 23:29 Tvir5.599999904632568-zeta19.03700065612793-N20000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 30 23:11 Tvir4.400000095367432-zeta131.34100341796875-N20000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 30 20:08 Tvir4.800000190734863-zeta131.34100341796875-N15000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 30 19:50 Tvir5.4770002365112305-zeta200.0-N15000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 30 19:32 Tvir4.698999881744385-zeta30.0-N15000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 30 19:14 Tvir5.599999904632568-zeta19.03700065612793-N15000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 30 18:57 Tvir4.400000095367432-zeta131.34100341796875-N15000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 12:41 Tvir4.800000190734863-zeta131.34100341796875-N7000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 12:23 Tvir5.4770002365112305-zeta200.0-N7000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 12:06 Tvir4.698999881744385-zeta30.0-N7000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 11:48 Tvir5.599999904632568-zeta19.03700065612793-N7000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 11:30 Tvir4.400000095367432-zeta131.34100341796875-N7000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 04:56 Tvir4.800000190734863-zeta131.34100341796875-N25600.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 04:38 Tvir5.4770002365112305-zeta200.0-N25600.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 04:21 Tvir4.698999881744385-zeta30.0-N25600.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 04:03 Tvir5.599999904632568-zeta19.03700065612793-N25600.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 03:45 Tvir4.400000095367432-zeta131.34100341796875-N25600.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 00:35 Tvir4.800000190734863-zeta131.34100341796875-N3000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 29 00:17 Tvir5.4770002365112305-zeta200.0-N3000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 23:59 Tvir4.698999881744385-zeta30.0-N3000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 23:42 Tvir5.599999904632568-zeta19.03700065612793-N3000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 23:20 Tvir4.400000095367432-zeta131.34100341796875-N3000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 21:06 Tvir4.800000190734863-zeta131.34100341796875-N10000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 20:49 Tvir5.4770002365112305-zeta200.0-N10000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 20:31 Tvir4.698999881744385-zeta30.0-N10000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 20:13 Tvir5.599999904632568-zeta19.03700065612793-N10000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 19:56 Tvir4.400000095367432-zeta131.34100341796875-N10000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 18:30 Tvir4.800000190734863-zeta131.34100341796875-N1000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 18:13 Tvir5.4770002365112305-zeta200.0-N1000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 17:55 Tvir4.698999881744385-zeta30.0-N1000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 17:37 Tvir5.599999904632568-zeta19.03700065612793-N1000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 17:20 Tvir4.400000095367432-zeta131.34100341796875-N1000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 28 14:03 Tvir4.400000095367432-zeta131.34100341796875-N5000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 10 18:58 Tvir4.800000190734863-zeta131.34100341796875-N5000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 10 18:40 Tvir5.4770002365112305-zeta200.0-N5000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 10 18:22 Tvir4.698999881744385-zeta30.0-N5000.npy\n",
+      "-rw-r--r--  1 bxia34 3.1M Jun 10 18:05 Tvir5.599999904632568-zeta19.03700065612793-N5000.npy\n"
      ]
-    },
+    }
+   ],
+   "source": [
+    "ll -lth outputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "params loaded: (3200, 2)\n",
-      "images rescaled to [-1.0, 1.1910929679870605]\n",
-      "params rescaled to [0.0002077208516410245, 0.999968750248265]\n"
+      "Number of parameters for nn_model: 111048705\n",
+      "sampling 800 images with normalized params = tensor([[0.2000, 0.5056]])\n",
+      "nn_model resumed from ./outputs/model_state-N3200\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "356d54e499ff46548303ff8d0a0b9520",
+       "model_id": "2d3427d677774c9785ee081b3b3b5542",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.8000, 0.0377]])\n",
+      "nn_model resumed from ./outputs/model_state-N3200\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "86370fad142545bbb095427ae255d9b2",
+       "model_id": "8f30542543bf4d96ac6284da1d3e2d91",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.3495, 0.0833]])\n",
+      "nn_model resumed from ./outputs/model_state-N3200\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a57710f5d32f4f1e8b29a7f5f86ee765",
+       "model_id": "8b55d00d4ec74b1995fabcb27152a20c",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.7385, 0.7917]])\n",
+      "nn_model resumed from ./outputs/model_state-N3200\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6c0df4be123e478b90e2b72e4ee9861a",
+       "model_id": "3fbfb9641b7c4709ab24f843b9ef9a41",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.4000, 0.5056]])\n",
+      "nn_model resumed from ./outputs/model_state-N3200\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "81f42c0247644f5692c588d0f1449b89",
+       "model_id": "f146207ea2af4e2fbbe19a7255fe13bd",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of parameters for nn_model: 111048705\n",
+      "sampling 800 images with normalized params = tensor([[0.2000, 0.5056]])\n",
+      "nn_model resumed from ./outputs/model_state-N6400\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fc91ff958e59444fac853802648e434b",
+       "model_id": "932fde5fc32e46719e809370a0145171",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.8000, 0.0377]])\n",
+      "nn_model resumed from ./outputs/model_state-N6400\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "291e28e5410f4fcc83b2b94e7b9f0183",
+       "model_id": "ad63275c7ba7499aa34c7b7a1d600b01",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.3495, 0.0833]])\n",
+      "nn_model resumed from ./outputs/model_state-N6400\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "540e85cd21484ec5bcdfd189af680414",
+       "model_id": "b92c97678a144c7db39e7504d9b51fdd",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.7385, 0.7917]])\n",
+      "nn_model resumed from ./outputs/model_state-N6400\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "56d6a74be233405cbd6c1c815775edf9",
+       "model_id": "ce95c55a33e94da39ad4761af0bca5da",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.4000, 0.5056]])\n",
+      "nn_model resumed from ./outputs/model_state-N6400\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "40d3fcf4469941c58c314989390a8037",
+       "model_id": "c7d30d6d4e9c46d6930da9b6d10e53c1",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of parameters for nn_model: 111048705\n",
+      "sampling 800 images with normalized params = tensor([[0.2000, 0.5056]])\n",
+      "nn_model resumed from ./outputs/model_state-N12800\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6e47ab612cc6440881f69245c985dbe5",
+       "model_id": "205b299c574343fca3f0f7bcce9c0fda",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.8000, 0.0377]])\n",
+      "nn_model resumed from ./outputs/model_state-N12800\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6b0771ada35043ed903f8535a0614ecd",
+       "model_id": "8e2fea1d73a54d87984bc430584ed7ee",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.3495, 0.0833]])\n",
+      "nn_model resumed from ./outputs/model_state-N12800\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a8b074d6e0734261b044f3420f14fc21",
+       "model_id": "8a4762d957a04a81b6283e303a0c1aa3",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.7385, 0.7917]])\n",
+      "nn_model resumed from ./outputs/model_state-N12800\n"
+     ]
+    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d0c0253596c643ed90ce22c5699312ed",
+       "model_id": "67dabb9f208d4977bdb8180afc4eec20",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "93ce72e5832241e480a21b58f8d9aba9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3500835f1f94425e8ee1eecfb2587934",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "da8089a140354180bef9132ed583938d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f4b9471044a34ccd8ea3894bc21e1b2f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3f771c1822f34d848da17e797eedfd71",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "89f214290fcf470b9d50dad8d943cf68",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "73f71bfcbb504de8a84d9308c2811bdb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "08ae268e1da244e0af842e70128ae154",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d2768d60e25b43f9b3307c55d10982d9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "77e3b9b7097142e2bd29e9004e774a9e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "aea446f4f98948e894b0cf42b66334c9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "12317786c96049cea7e590cdcf6a504f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "89d95608fc7a4dc1b4130f25a77a0562",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f1db17ea59134aa58a1386388905a674",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "649a0196def7425c98657a0f0f7cf331",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "851b430a32e64399aafd472bf8d5db20",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "080379d2e3804d71a063f4dbeb29adbf",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7083e05978cc427e8f8a6d7af7e40124",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7424d068a41b4331ba77a3296b719403",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ac141a1cc1ad4cb19f245c1409b9ea55",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fd456c7b00064808a9a710b921110492",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2f2cb8b2c64d437b9be754fdcac1f2c8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "14a78d1597214b03aa09d947853673bb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c05038e3d751431d940d04c9f886c6b6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "abe0bceafda84155b8637721e2d97fcc",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1852a512bc684aacb47d66b95bcd229b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "09d9c8b9e18e4475b69b258262c58976",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "65598ff9820a42c691816ba532a81cdb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f3205176ebd745b3813e218acfedb77d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "30332b0311d24babb5874067a6f9b699",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "95eda7879d1d4e6bb1c6085085cf41c0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9b1abdfb85b64b688e0d20963a5a021d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e559649999b34ac6911b4377c20cec44",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "414df27d9b4e4b08962c15f8156f2199",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a029bc6c3fe74c009143efed7ab79134",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "754b5dbd719f4d66aaab63d6160fa40f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/64 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "saved model at ./outputs/model_state-N3200\n",
-      "Number of parameters for nn_model: 111048705\n",
-      "---------------- num_image = 6400 ----------------\n",
-      "run_name = 0702-1737\n",
-      "Launching training on one GPU.\n",
-      "dataset content: <KeysViewHDF5 ['brightness_temp', 'density', 'kwargs', 'params', 'redshifts_distances', 'seeds', 'xH_box']>\n",
-      "51200 images can be loaded\n",
-      "field.shape = (64, 64, 514)\n",
-      "params keys = [b'ION_Tvir_MIN', b'HII_EFF_FACTOR']\n",
-      "loading 6400 images randomly\n",
-      "images loaded: (6400, 1, 64, 64)\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Detected kernel version 3.10.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "params loaded: (6400, 2)\n",
-      "images rescaled to [-1.0, 1.1910929679870605]\n",
-      "params rescaled to [0.0001702067256199591, 0.9998461462686923]\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7dcbe3c1221443a0910da2952ba6b58d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5f092d9c36bc4e9ebe8fb91255c477cf",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6928d46e65b94dcb9ba1b52d919b438b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5a6bea10bb534ebbb3308a3feba3d2f5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "38d144c4047841e4b151ceba3ea04e69",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b66572a96d874696bb5019f6c3c943b0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4ff636158a724298ba77fe9de939cbb1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4d64ae27ef0546f581270fabc1ce1bb5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a38cfc5df2c245b6b7039ce0f1e7377e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fa6228cdfed14eb29d354fefe892b9a7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "835f54dfb7d746f1ab7a215568dbcd6b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "30a4aef103d945b0bbf2bd2b442d8c92",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c15ff3f29f0e41edb871ded32b6e21e2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a34ba7e833194118a39af751a56c13f8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ca66fa6562544d3b977138955e555278",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a76d013ffe314ac1af6ef055f392415d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e43c637e27fa46919a4ea48aed02a1d0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2508f4acc53840559da06118764964c2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "52e7c4532c4a4025bb1d9b20c0d9f45b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4cf2e1ac6418485cade29e15f4514684",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "00111e2eb38c4d0ca54c20345feff058",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0ac2eda548164d93b45827112252a16a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "521967e01f6042bfbeb284518198cc2b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4bc1af553d4e4ae188620d4eda8d9721",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c63a5f2eeef64e08997b05cb4343ca60",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c765aaa74ac048cf8f30d266c35dbd08",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "147fd29ff61647478642470c449b0bd0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "548617e45b924dc99b0294858eda91f1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "29a200c7fe2d46cdb790f0481c10b652",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "eaa1e415c8794b99be33795756e51696",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ab3111ab9dcc44209ed83acd718bb833",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e9d385fea7014bdfb1273f5b582ee675",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "83f012b3b2644fa3a587c09993bf64ab",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ef0be31575214e7091bdc6d9db69d012",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "518509bc011e4d4e8dbdab22416b4077",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "58e281c972d14d9eb1227c3c29347804",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "48922d19888b45e39e421c23a7271ed2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e4920eb2e56842eeafe7d5e687179811",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "df660d2b843a44dd82b8b3c02c928a12",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1de90798d93b42ccb09e48fb9e78abc5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c92a406b0fa644d5b4cf010db1dcb7e3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "34b25fe26b4646efb0efbeaf01bd4eaf",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "45c9a02919ec4fb2abf2cf70ba36a815",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1d17bde3d2db4aee8608a907e8e891d8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "94e48fce50434a698720af7cd694d429",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cac6913acbf74365be12d902846e9894",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "aa8f2c48117849e2a8d146b1397cfedb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8ebe2d4e22a345af97973e20af55f876",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b72a5693076a405bbbb63020c51dfec9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "91b2c17142df40bcac99e8933052f2e5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/128 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "saved model at ./outputs/model_state-N6400\n",
-      "Number of parameters for nn_model: 111048705\n",
-      "--------------- num_image = 12800 ----------------\n",
-      "run_name = 0702-1801\n",
-      "Launching training on one GPU.\n",
-      "dataset content: <KeysViewHDF5 ['brightness_temp', 'density', 'kwargs', 'params', 'redshifts_distances', 'seeds', 'xH_box']>\n",
-      "51200 images can be loaded\n",
-      "field.shape = (64, 64, 514)\n",
-      "params keys = [b'ION_Tvir_MIN', b'HII_EFF_FACTOR']\n",
-      "loading 12800 images randomly\n",
-      "images loaded: (12800, 1, 64, 64)\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Detected kernel version 3.10.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "params loaded: (12800, 2)\n",
-      "images rescaled to [-1.0, 1.2380504608154297]\n",
-      "params rescaled to [3.8339325512049e-05, 0.9999958400767995]\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7a44e73529d34fcdbc790b622ce0af6a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c04f1c7c5ba04e4388bccba5e2996d6b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "24c1aab784f74750a4484bff59df2a3f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "de7bde4f1e0c4fb89f29b3fd4ec45be7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a89dbe6280ab46e4a2e59eaabb214cbe",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6c70e9de3be3499ea65973ebadb68dbb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bf561aefab6b4247a25aa1801d0d73c9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fbd31dfa0e384639ac6662fc937df56c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8e144dd4e2c240dc829e1f6c319cd9ac",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "733e8a15f5154eb8999cb567c401df77",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "701d2720e78b48658bec374f8b600a96",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cfac805267624d2f850a6bc85bbd565e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ccbd833ee5ac4e7e9ea689f75a559999",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "492bb0481fd54207a54881b4cf963ab7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1c10e32bce134ad7a374576e0a912a4d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7255a31d029c477aa780ab8146152d33",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "044db3abab9d4842ab706bedd228dec9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "733e11e6b81148419dc48bc7933a1b30",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7952f16d57b94846be703767116e05f3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a38997af95174caf893abb93dcdbad6a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "62123abb03c54098b46c3d5a5d43284e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8c0838ab3635462d992a9d630700130e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "afeddb3d0a7a42b69a1cfa2ada68ea5f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1fc0135674e84e7eafea406c737f43f3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d6228335f2d54c06804583a2caf44a57",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f040934616984913b01328d419291b31",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "12c2e4009ec94c20a7bc7401a11d344a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f50c74d91d4047e6995c18948cd77c38",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "65562fe29dc544678bd5825f681f2066",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9b55fbb6c10247c5a5cde74a88ba39d1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c093b9c1819c4541a0dcbd92787ceef3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1778af8b0e7e43e5b8f702b53b4526e6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "34589991352a4bf3891ba33ebfcd521b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3acb4d18d1c94cc4930804fd1501589d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f0a6fcb3493c4f00a5ba4bf9205fdd83",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dd59413239e74b2fa0b2ebe7b9eb0e50",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8ffa4b42bea149e3a940c1fb492ebf1c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "42091af8488e4004a9078590e48db390",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "00b4ab56b15f408d81796a9d143867cd",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "355be85a9a3a490c8eebe7e4bf97ae51",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6d6d22b44e774a4b9af42a469b6e33c0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "53c981f957504a2999c8677e9dd26775",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "17eb6a178bda4e52b4257e75de4d388d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c121efbc28ec4ae894bc220ba12c4cd3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8a1f2197cea8456992b7b9df83a108ba",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "71af44b6a67d41de949b8675734c85b8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "396dbc8bfb684400bfa23708ac610e30",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1212909241094340bd81240b1565fa50",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d3fbfded55e44462be7ebdcd1850b3ad",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d27e62debe9f496bbed45cf4f803cda8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/256 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "saved model at ./outputs/model_state-N12800\n",
-      "Number of parameters for nn_model: 111048705\n",
-      "--------------- num_image = 25600 ----------------\n",
-      "run_name = 0702-1845\n",
-      "Launching training on one GPU.\n",
-      "dataset content: <KeysViewHDF5 ['brightness_temp', 'density', 'kwargs', 'params', 'redshifts_distances', 'seeds', 'xH_box']>\n",
-      "51200 images can be loaded\n",
-      "field.shape = (64, 64, 514)\n",
-      "params keys = [b'ION_Tvir_MIN', b'HII_EFF_FACTOR']\n",
-      "loading 25600 images randomly\n",
-      "images loaded: (25600, 1, 64, 64)\n",
-      "params loaded: (25600, 2)\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Detected kernel version 3.10.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "images rescaled to [-1.0, 1.2380504608154297]\n",
-      "params rescaled to [3.349517406547875e-06, 0.9999922179553216]\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "488a6a64228f4ca487199a34fe14f334",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dd65d00b70b548138c47ac135d34d885",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fc04006857674f30b1cd269474ec2951",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f6489109774c44cd9028f5f3e2a43c78",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1b740f503d21482bb328e6008f36309f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "49c3a028867b46399556d72c87ab4b91",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "be6bdf8c976d424fb2b7b84cb4a3126e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "414889f7a16344fead807080574c135f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ce71461129bc494d998d20ea698d0c6d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0f3e74076ea54a03b65a3a22d00a7e39",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0cb0360d7c864b6abe628c19f7f767f1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b601b32a3845432f9cef4957fb1f2409",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "55f3c7ef38c942aaab603a0672c2b119",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6cbf0716b386442fbf605817fd69bc42",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "752a28c480314f33bfdb140fd20b1fac",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "35c111a90566421fb8b84d2111bd1b02",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3def96bb7d9c4536b83cffd259d7dcd3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fb0f3e6442d146489219914f1eef83b5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "97a889bf4a5b48ccbc608d4c3cca9271",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4f9405081dae4e1293ab3d2759cfb7e7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "665e98865f764ca68daaf450b1be5298",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "26311c1f6dc24236b41a4f39a216fa77",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "67c142764fd044ed94b072daaa3c8c4e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "41ecea15ccbb44a9967905aad3942436",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "14f4a12e186144309e9cc0acb5a4f129",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d1ea77df2a754fff9d9757fa03a9f7e0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8a317145dc034be38138af1a01a9d934",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "550ef8709fa14a6ba4ecdaf8f3d216d9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "790532f8da474ec3a5c79b998e31afce",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2c8c5887cd604cf89793125301208673",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2b801564cef740828b6c7c42a514d1a5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1761b1f1968f4fa3a7f62973c4098523",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "19bde9c1086f4133abbe4209a81a7c07",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "98c328ca977847b6a017515563657c5e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "55a1e6800b4d49b28a2a10ffcf61c7cb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9a8ea0ee7296474aa4a3c872c9195906",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "713c51e58a9842f680e24f4b8c0ed3e2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "98ced291ac0a4c9db980de88e3a60052",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d74c67acd7cb4e589a3a0c2c131e56e4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fcd50d883806450db2dd6423f2515898",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bac9b4a14afb4abe948f27222d93f475",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "83fcb65fc5ac46788b2b7a10e5b2f93c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6e231e56299f4b0db7cb374ef6a9fe24",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "53194afa16ea46fdbf1762dc80d7b311",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.4000, 0.5056]])\n",
+      "nn_model resumed from ./outputs/model_state-N12800\n"
+     ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3e43bde8d11c499aa6a4bd359838fac5",
+       "model_id": "950ba8b752444e1e9a2d817968f5c3a3",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1fc41813946b41da92945fae500ca478",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Number of parameters for nn_model: 111048705\n",
+      "sampling 800 images with normalized params = tensor([[0.2000, 0.5056]])\n",
+      "nn_model resumed from ./outputs/model_state-N25600\n"
+     ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "176c57c2d1d24c59b65c850e47b40c2f",
+       "model_id": "fd753c01957447398ec61ff383616825",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "04424a5f8e7d4e97ae32ce3f8eec89d9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sampling 800 images with normalized params = tensor([[0.8000, 0.0377]])\n",
+      "nn_model resumed from ./outputs/model_state-N25600\n"
+     ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "03f41af225f642e88c00d805d8bcef20",
+       "model_id": "a328cf05873a4370a645270ea917233d",
        "version_major": 2,
        "version_minor": 0
       },
       "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
+       "  0%|          | 0/1000 [00:00<?, ?it/s]"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a9c200264c294211b0435fad26c474fe",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "  0%|          | 0/512 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "if __name__ == \"__main__\":\n",
-    "    # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)\n",
-    "    config = TrainConfig()\n",
-    "    for i, num_image in enumerate(num_image_list):\n",
-    "        config.num_image = num_image\n",
-    "        ddpm21cm = DDPM21CM(config)\n",
-    "        print(f\" num_image = {ddpm21cm.config.num_image} \".center(50, '-'))\n",
-    "        print(f\"run_name = {ddpm21cm.config.run_name}\")\n",
-    "        notebook_launcher(ddpm21cm.train, num_processes=1)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "total 1.8G\n",
-      "-rw-r--r--  1 bxia34 848M Jul  2 16:31 model_state-N3200\n",
-      "drwxr-xr-x 12 bxia34 4.0K Jul  2 16:31 \u001b[0m\u001b[01;34mlogs\u001b[0m/\n",
-      "-rw-r--r--  1 bxia34 848M Jul  2 16:18 model_state-N1600\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 12:31 Tvir4.800000190734863-zeta131.34100341796875-N2000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 12:12 Tvir5.4770002365112305-zeta200.0-N2000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 11:54 Tvir4.698999881744385-zeta30.0-N2000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 11:35 Tvir5.599999904632568-zeta19.03700065612793-N2000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 11:17 Tvir4.400000095367432-zeta131.34100341796875-N2000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 10:37 Tvir4.800000190734863-zeta131.34100341796875-N32000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 04:47 Tvir5.4770002365112305-zeta200.0-N32000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 04:29 Tvir4.698999881744385-zeta30.0-N32000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 04:11 Tvir5.599999904632568-zeta19.03700065612793-N32000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 03:53 Tvir4.400000095367432-zeta131.34100341796875-N32000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 00:23 Tvir4.800000190734863-zeta131.34100341796875-N20000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jul  1 00:05 Tvir5.4770002365112305-zeta200.0-N20000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 30 23:47 Tvir4.698999881744385-zeta30.0-N20000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 30 23:29 Tvir5.599999904632568-zeta19.03700065612793-N20000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 30 23:11 Tvir4.400000095367432-zeta131.34100341796875-N20000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 30 20:08 Tvir4.800000190734863-zeta131.34100341796875-N15000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 30 19:50 Tvir5.4770002365112305-zeta200.0-N15000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 30 19:32 Tvir4.698999881744385-zeta30.0-N15000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 30 19:14 Tvir5.599999904632568-zeta19.03700065612793-N15000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 30 18:57 Tvir4.400000095367432-zeta131.34100341796875-N15000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 12:41 Tvir4.800000190734863-zeta131.34100341796875-N7000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 12:23 Tvir5.4770002365112305-zeta200.0-N7000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 12:06 Tvir4.698999881744385-zeta30.0-N7000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 11:48 Tvir5.599999904632568-zeta19.03700065612793-N7000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 11:30 Tvir4.400000095367432-zeta131.34100341796875-N7000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 04:56 Tvir4.800000190734863-zeta131.34100341796875-N25600.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 04:38 Tvir5.4770002365112305-zeta200.0-N25600.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 04:21 Tvir4.698999881744385-zeta30.0-N25600.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 04:03 Tvir5.599999904632568-zeta19.03700065612793-N25600.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 03:45 Tvir4.400000095367432-zeta131.34100341796875-N25600.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 00:35 Tvir4.800000190734863-zeta131.34100341796875-N3000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 29 00:17 Tvir5.4770002365112305-zeta200.0-N3000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 23:59 Tvir4.698999881744385-zeta30.0-N3000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 23:42 Tvir5.599999904632568-zeta19.03700065612793-N3000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 23:20 Tvir4.400000095367432-zeta131.34100341796875-N3000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 21:06 Tvir4.800000190734863-zeta131.34100341796875-N10000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 20:49 Tvir5.4770002365112305-zeta200.0-N10000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 20:31 Tvir4.698999881744385-zeta30.0-N10000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 20:13 Tvir5.599999904632568-zeta19.03700065612793-N10000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 19:56 Tvir4.400000095367432-zeta131.34100341796875-N10000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 18:30 Tvir4.800000190734863-zeta131.34100341796875-N1000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 18:13 Tvir5.4770002365112305-zeta200.0-N1000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 17:55 Tvir4.698999881744385-zeta30.0-N1000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 17:37 Tvir5.599999904632568-zeta19.03700065612793-N1000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 17:20 Tvir4.400000095367432-zeta131.34100341796875-N1000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 28 14:03 Tvir4.400000095367432-zeta131.34100341796875-N5000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 10 18:58 Tvir4.800000190734863-zeta131.34100341796875-N5000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 10 18:40 Tvir5.4770002365112305-zeta200.0-N5000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 10 18:22 Tvir4.698999881744385-zeta30.0-N5000.npy\n",
-      "-rw-r--r--  1 bxia34 3.1M Jun 10 18:05 Tvir5.599999904632568-zeta19.03700065612793-N5000.npy\n"
-     ]
-    }
-   ],
-   "source": [
-    "ll -lth outputs"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Number of parameters for nn_model: 111048705\n",
-      "sampling 800 images with normalized params = tensor([[0.2000, 0.5056]])\n",
-      "nn_model resumed from ./outputs/model_state-N1600\n"
+      "sampling 800 images with normalized params = tensor([[0.3495, 0.0833]])\n",
+      "nn_model resumed from ./outputs/model_state-N25600\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1a2676fba899451a83430c7884b20f2d",
+       "model_id": "3838c0ba114c406caa3e39668feaea38",
        "version_major": 2,
        "version_minor": 0
       },
@@ -4346,14 +1799,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "sampling 800 images with normalized params = tensor([[0.8000, 0.0377]])\n",
-      "nn_model resumed from ./outputs/model_state-N1600\n"
+      "sampling 800 images with normalized params = tensor([[0.7385, 0.7917]])\n",
+      "nn_model resumed from ./outputs/model_state-N25600\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "234a1884d99b4f838d42290e1dad2755",
+       "model_id": "a668e8fd474e4d659a8960e35598fb5f",
        "version_major": 2,
        "version_minor": 0
       },
@@ -4368,15 +1821,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Number of parameters for nn_model: 111048705\n",
-      "sampling 800 images with normalized params = tensor([[0.2000, 0.5056]])\n",
-      "nn_model resumed from ./outputs/model_state-N3200\n"
+      "sampling 800 images with normalized params = tensor([[0.4000, 0.5056]])\n",
+      "nn_model resumed from ./outputs/model_state-N25600\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "606fa802d13546d9b643cf4c033aeb3a",
+       "model_id": "028b3cc2a3214999b6a76700579c4263",
        "version_major": 2,
        "version_minor": 0
       },
@@ -4386,24 +1838,12 @@
      },
      "metadata": {},
      "output_type": "display_data"
-    },
-    {
-     "ename": "KeyboardInterrupt",
-     "evalue": "",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[11], line 9\u001b[0m\n\u001b[1;32m      6\u001b[0m config\u001b[39m.\u001b[39mnum_image \u001b[39m=\u001b[39m num_image\n\u001b[1;32m      7\u001b[0m ddpm21cm \u001b[39m=\u001b[39m DDPM21CM(config)\n\u001b[0;32m----> 9\u001b[0m ddpm21cm\u001b[39m.\u001b[39;49msample(\u001b[39mf\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39m./outputs/model_state-N\u001b[39;49m\u001b[39m{\u001b[39;49;00mnum_image\u001b[39m}\u001b[39;49;00m\u001b[39m\"\u001b[39;49m, params\u001b[39m=\u001b[39;49mtorch\u001b[39m.\u001b[39;49mtensor([\u001b[39m4.4\u001b[39;49m, \u001b[39m131.341\u001b[39;49m]), repeat\u001b[39m=\u001b[39;49mrepeat)\n\u001b[1;32m     11\u001b[0m ddpm21cm\u001b[39m.\u001b[39msample(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m./outputs/model_state-N\u001b[39m\u001b[39m{\u001b[39;00mnum_image\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m, params\u001b[39m=\u001b[39mtorch\u001b[39m.\u001b[39mtensor((\u001b[39m5.6\u001b[39m, \u001b[39m19.037\u001b[39m)), repeat\u001b[39m=\u001b[39mrepeat)\n\u001b[1;32m     13\u001b[0m \u001b[39m# ddpm21cm.sample(f\"./outputs/model_state-N{num_image}\", params=torch.tensor((4.699, 30)), repeat=repeat)\u001b[39;00m\n\u001b[1;32m     14\u001b[0m \n\u001b[1;32m     15\u001b[0m \u001b[39m# ddpm21cm.sample(f\"./outputs/model_state-N{num_image}\", params=torch.tensor((5.477, 200)), repeat=repeat)\u001b[39;00m\n\u001b[1;32m     16\u001b[0m \n\u001b[1;32m     17\u001b[0m \u001b[39m# ddpm21cm.sample(f\"./outputs/model_state-N{num_image}\", params=torch.tensor((4.8, 131.341)), repeat=repeat)\u001b[39;00m\n",
-      "Cell \u001b[0;32mIn[5], line 207\u001b[0m, in \u001b[0;36mDDPM21CM.sample\u001b[0;34m(self, file, params, repeat, ema, entire)\u001b[0m\n\u001b[1;32m    202\u001b[0m \u001b[39m# self.ema_model = ContextUnet(n_param=config.n_param, image_size=config.HII_DIM, dim=config.dim, stride=config.stride).to(config.device)\u001b[39;00m\n\u001b[1;32m    203\u001b[0m \u001b[39m# self.ema_model.load_state_dict(torch.load(os.path.join(config.output_dir, f\"{config.resume}\"))['ema_unet_state_dict'])\u001b[39;00m\n\u001b[1;32m    204\u001b[0m \u001b[39m# print(f\"resumed ema_model from {config.resume}\")\u001b[39;00m\n\u001b[1;32m    206\u001b[0m \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39mno_grad():\n\u001b[0;32m--> 207\u001b[0m     x_last, x_entire \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mddpm\u001b[39m.\u001b[39;49msample(\n\u001b[1;32m    208\u001b[0m         nn_model\u001b[39m=\u001b[39;49mnn_model, \n\u001b[1;32m    209\u001b[0m         params\u001b[39m=\u001b[39;49mparams\u001b[39m.\u001b[39;49mto(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconfig\u001b[39m.\u001b[39;49mdevice), \n\u001b[1;32m    210\u001b[0m         device\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconfig\u001b[39m.\u001b[39;49mdevice, \n\u001b[1;32m    211\u001b[0m         guide_w\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconfig\u001b[39m.\u001b[39;49mguide_w\n\u001b[1;32m    212\u001b[0m         )\n\u001b[1;32m    214\u001b[0m \u001b[39m# np.save(os.path.join(self.config.output_dir, f\"{self.config.run_name}{'ema' if ema else ''}.npy\"), x_last)\u001b[39;00m\n\u001b[1;32m    215\u001b[0m np\u001b[39m.\u001b[39msave(os\u001b[39m.\u001b[39mpath\u001b[39m.\u001b[39mjoin(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconfig\u001b[39m.\u001b[39moutput_dir, \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mTvir\u001b[39m\u001b[39m{\u001b[39;00mparams_backup[\u001b[39m0\u001b[39m]\u001b[39m}\u001b[39;00m\u001b[39m-zeta\u001b[39m\u001b[39m{\u001b[39;00mparams_backup[\u001b[39m1\u001b[39m]\u001b[39m}\u001b[39;00m\u001b[39m-N\u001b[39m\u001b[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconfig\u001b[39m.\u001b[39mnum_image\u001b[39m}\u001b[39;00m\u001b[39m{\u001b[39;00m\u001b[39m'\u001b[39m\u001b[39mema\u001b[39m\u001b[39m'\u001b[39m\u001b[39m \u001b[39m\u001b[39mif\u001b[39;00m\u001b[39m \u001b[39mema\u001b[39m \u001b[39m\u001b[39melse\u001b[39;00m\u001b[39m \u001b[39m\u001b[39m'\u001b[39m\u001b[39m'\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39m.npy\u001b[39m\u001b[39m\"\u001b[39m), x_last)\n",
-      "Cell \u001b[0;32mIn[2], line 59\u001b[0m, in \u001b[0;36mDDPMScheduler.sample\u001b[0;34m(self, nn_model, params, device, guide_w)\u001b[0m\n\u001b[1;32m     56\u001b[0m pbar_sample\u001b[39m.\u001b[39mset_description(\u001b[39m\"\u001b[39m\u001b[39mSampling\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[1;32m     57\u001b[0m \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m \u001b[39mreversed\u001b[39m(\u001b[39mrange\u001b[39m(\u001b[39m0\u001b[39m, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnum_timesteps)):\n\u001b[1;32m     58\u001b[0m     \u001b[39m# print(f'sampling timestep {i:4d}',end='\\r')\u001b[39;00m\n\u001b[0;32m---> 59\u001b[0m     t_is \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39;49mtensor([i])\u001b[39m.\u001b[39;49mto(device)\n\u001b[1;32m     60\u001b[0m     t_is \u001b[39m=\u001b[39m t_is\u001b[39m.\u001b[39mrepeat(n_sample)\n\u001b[1;32m     62\u001b[0m     z \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mrandn(n_sample, \u001b[39m*\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mimg_shape)\u001b[39m.\u001b[39mto(device) \u001b[39mif\u001b[39;00m i \u001b[39m>\u001b[39m \u001b[39m0\u001b[39m \u001b[39melse\u001b[39;00m \u001b[39m0\u001b[39m\n",
-      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
-     ]
     }
    ],
    "source": [
     "if __name__ == \"__main__\":\n",
     "    num_image_list = [1600,3200,6400,12800,25600]\n",
+    "    # num_image_list = [3200,6400,12800,25600]\n",
     "    # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)\n",
     "    repeat = 800\n",
     "    config = TrainConfig()\n",