0522-0031

diffusion.ipynb

@@ -32,7 +32,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f925cb378800455fb1216e84a2900e58",
+       "model_id": "b9a1289a64b14b6e9aebf8148c38c9a8",
        "version_major": 2,
        "version_minor": 0
       },
@@ -234,7 +234,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 84,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -278,7 +278,7 @@
     "    def sample(self, nn_model, params, device, guide_w = 0):\n",
     "        n_sample = len(params) #params.shape[0]\n",
     "        # print(\"params.shape[0], len(params)\", params.shape[0], len(params))\n",
-    "        x_i = torch.randn(n_sample, *self.img_shape[1:]).to(device)\n",
+    "        x_i = torch.randn(n_sample, *self.img_shape).to(device)\n",
     "        print(\"x_i.shape =\", x_i.shape)\n",
     "        # print(\"x_i.shape =\", x_i.shape)\n",
     "        if guide_w != -1:\n",
@@ -299,7 +299,7 @@
     "            t_is = torch.tensor([i]).to(device)\n",
     "            t_is = t_is.repeat(n_sample)\n",
     "\n",
-    "            z = torch.randn(n_sample, *self.img_shape[1:]).to(device) if i > 0 else 0\n",
+    "            z = torch.randn(n_sample, *self.img_shape).to(device) if i > 0 else 0\n",
     "\n",
     "            if guide_w == -1:\n",
     "                # eps = nn_model(x_i, t_is, return_dict=False)[0]\n",
@@ -307,7 +307,7 @@
     "                # x_i = 1/torch.sqrt(self.alpha_t[i])*(x_i-eps*self.beta_t[i]/torch.sqrt(1-self.bar_alpha_t[i])) + torch.sqrt(self.beta_t[i])*z\n",
     "            else:\n",
     "                # double batch\n",
-    "                x_i = x_i.repeat(2, *torch.ones(len(self.img_shape[1:]), dtype=int).tolist())\n",
+    "                x_i = x_i.repeat(2, *torch.ones(len(self.img_shape), dtype=int).tolist())\n",
     "                t_is = t_is.repeat(2)\n",
     "\n",
     "                # split predictions and compute weighting\n",
@@ -341,7 +341,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 85,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -383,7 +383,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 86,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -408,7 +408,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 87,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -438,7 +438,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 88,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -453,7 +453,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 89,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -467,7 +467,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 90,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -485,7 +485,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 91,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -567,7 +567,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 92,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -600,7 +600,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 93,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -649,7 +649,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 94,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -678,7 +678,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 95,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -934,7 +934,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 96,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -964,7 +964,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 100,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1030,7 +1030,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 101,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1040,7 +1040,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 102,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1049,7 +1049,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 103,
+   "execution_count": 23,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1073,7 +1073,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 104,
+   "execution_count": 24,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1271,7 +1271,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 105,
+   "execution_count": 25,
    "metadata": {},
    "outputs": [
     {
@@ -1429,10 +1429,11 @@
     "                        print('saved model at ' + self.config.output_dir + f\"model_state_{ep:02d}.pth\")\n",
     "                        # print('saved model at ' + config.save_dir + f\"model_epoch_{ep}_test_{config.run_name}.pth\")\n",
     "\n",
-    "    def sample(self, file, params=torch.tensor((0.2,0.8)).view(1,2), ema=False, entire=False):\n",
+    "    def sample(self, file, params=None, ema=False, entire=False):\n",
     "        # n_sample = params.shape[0]\n",
+    "        params = params or torch.tensor([0.2,0.8]).repeat(10,1)\n",
     "        assert params.dim() == 2, \"params must be a 2D torch.tensor\"\n",
-    "        # print(\"params.shape =\", params.shape)\n",
+    "        print(\"params.shape =\", params.shape)\n",
     "        # print(\"len(params) =\", len(params))\n",
     "        model = self.ema_model if ema else self.nn_model\n",
     "        # params = torch.tile(params, (n_sample,1)).to(device)\n",
@@ -1450,7 +1451,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 106,
+   "execution_count": 26,
    "metadata": {},
    "outputs": [
     {
@@ -1477,14 +1478,14 @@
      "output_type": "stream",
      "text": [
       "params loaded: (200, 2)\n",
-      "images rescaled to [-1.0, 1.1032476425170898]\n",
-      "params rescaled to [0.0, 0.9962284381407488]\n"
+      "images rescaled to [-1.0, 1.0288152694702148]\n",
+      "params rescaled to [0.0, 0.9958380936251952]\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "986e997d5f8c4c24b45d221987296db6",
+       "model_id": "fa014381477545e4a0fb6976d0cd3e0e",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1498,7 +1499,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2aab12179dd24a8788d8ce3dfcce6a34",
+       "model_id": "8bdd69fdc66c422c99b59032ef36b1f4",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1512,7 +1513,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1560051f3a3f460984173d8564715be7",
+       "model_id": "3cea3f68bad54e9bb4a63a897bea12dd",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1526,7 +1527,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d782dec46b654831a9bdd35482be746a",
+       "model_id": "26f29f7e73fd493fa62d52e6441589f1",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1540,7 +1541,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f633f6d189934559844de4a3bacfa8d4",
+       "model_id": "439448549c8840c0a970861a66e83264",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1554,7 +1555,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e4035b6e33df421e8c844839387f574b",
+       "model_id": "acb1905d443a451c937cc0c114e558c9",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1568,7 +1569,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "069d8cc887b04a4cbb2119342a0bf44e",
+       "model_id": "4f81671b7e1b43218b5aa6620e59503b",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1582,7 +1583,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1c58007af7184304a8bb2774236f9729",
+       "model_id": "66673bbcaffc412796022e1ed372b6cf",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1596,7 +1597,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b4af0c5b1460456488bea9de6e159a51",
+       "model_id": "2e04dfb1e9da4895a6f83516684bec03",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1610,7 +1611,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0483ca6f03fe4be094639c26fc660925",
+       "model_id": "1dc4006116ca46e0a2b74a74858fa57a",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1628,20 +1629,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 128,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "x_i.shape = torch.Size([1, 512])\n"
+      "params.shape = torch.Size([10, 2])\n",
+      "x_i.shape = torch.Size([10, 1, 64, 512])\n"
      ]
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cbef2a1163e74da5918a66fafc43caaa",
+       "model_id": "44830b578bb94dc5bf65eafaff19b1f5",
        "version_major": 2,
        "version_minor": 0
       },
@@ -1656,27 +1658,31 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "nn_model input shape torch.Size([2, 512]) torch.Size([2]) torch.Size([2, 2])\n"
+      "nn_model input shape torch.Size([20, 1, 64, 512]) torch.Size([20]) torch.Size([20, 2])\n"
      ]
     },
     {
      "ename": "RuntimeError",
-     "evalue": "Expected 3D (unbatched) or 4D (batched) input to conv2d, but got input of size: [2, 512]",
+     "evalue": "CUDA out of memory. Tried to allocate 320.00 MiB (GPU 0; 23.64 GiB total capacity; 21.71 GiB already allocated; 170.50 MiB free; 22.45 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF",
      "output_type": "error",
      "traceback": [
       "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
       "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[99], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m ddpm21cm\u001b[39m.\u001b[39;49msample(\u001b[39m\"\u001b[39;49m\u001b[39m./outputs/model_state_09.pth\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n",
-      "Cell \u001b[0;32mIn[98], line 153\u001b[0m, in \u001b[0;36mDDPM21CM.sample\u001b[0;34m(self, file, params, ema, entire)\u001b[0m\n\u001b[1;32m    150\u001b[0m model \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mema_model \u001b[39mif\u001b[39;00m ema \u001b[39melse\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnn_model\n\u001b[1;32m    151\u001b[0m \u001b[39m# params = torch.tile(params, (n_sample,1)).to(device)\u001b[39;00m\n\u001b[0;32m--> 153\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(model, 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), device\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconfig\u001b[39m.\u001b[39;49mdevice, 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    155\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[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconfig\u001b[39m.\u001b[39mrun_name\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[39mNone\u001b[39;00m\u001b[39m}\u001b[39;00m\u001b[39m.npy\u001b[39m\u001b[39m\"\u001b[39m), x_last)\n\u001b[1;32m    156\u001b[0m \u001b[39mif\u001b[39;00m entire:\n",
-      "Cell \u001b[0;32mIn[84], line 75\u001b[0m, in \u001b[0;36mDDPMScheduler.sample\u001b[0;34m(self, nn_model, params, device, guide_w)\u001b[0m\n\u001b[1;32m     73\u001b[0m \u001b[39m# split predictions and compute weighting\u001b[39;00m\n\u001b[1;32m     74\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mnn_model input shape\u001b[39m\u001b[39m\"\u001b[39m, x_i\u001b[39m.\u001b[39mshape, t_is\u001b[39m.\u001b[39mshape, c_i\u001b[39m.\u001b[39mshape)\n\u001b[0;32m---> 75\u001b[0m eps \u001b[39m=\u001b[39m nn_model(x_i, t_is, c_i)\n\u001b[1;32m     76\u001b[0m eps1 \u001b[39m=\u001b[39m eps[:n_sample]\n\u001b[1;32m     77\u001b[0m eps2 \u001b[39m=\u001b[39m eps[n_sample:]\n",
+      "Cell \u001b[0;32mIn[128], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m ddpm21cm\u001b[39m.\u001b[39;49msample(\u001b[39m\"\u001b[39;49m\u001b[39m./outputs/model_state_09.pth\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n",
+      "Cell \u001b[0;32mIn[127], line 154\u001b[0m, in \u001b[0;36mDDPM21CM.sample\u001b[0;34m(self, file, params, ema, entire)\u001b[0m\n\u001b[1;32m    151\u001b[0m model \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mema_model \u001b[39mif\u001b[39;00m ema \u001b[39melse\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mnn_model\n\u001b[1;32m    152\u001b[0m \u001b[39m# params = torch.tile(params, (n_sample,1)).to(device)\u001b[39;00m\n\u001b[0;32m--> 154\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(model, 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), device\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mconfig\u001b[39m.\u001b[39;49mdevice, 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    156\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[39m{\u001b[39;00m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconfig\u001b[39m.\u001b[39mrun_name\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[39mNone\u001b[39;00m\u001b[39m}\u001b[39;00m\u001b[39m.npy\u001b[39m\u001b[39m\"\u001b[39m), x_last)\n\u001b[1;32m    157\u001b[0m \u001b[39mif\u001b[39;00m entire:\n",
+      "Cell \u001b[0;32mIn[111], line 75\u001b[0m, in \u001b[0;36mDDPMScheduler.sample\u001b[0;34m(self, nn_model, params, device, guide_w)\u001b[0m\n\u001b[1;32m     73\u001b[0m \u001b[39m# split predictions and compute weighting\u001b[39;00m\n\u001b[1;32m     74\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mnn_model input shape\u001b[39m\u001b[39m\"\u001b[39m, x_i\u001b[39m.\u001b[39mshape, t_is\u001b[39m.\u001b[39mshape, c_i\u001b[39m.\u001b[39mshape)\n\u001b[0;32m---> 75\u001b[0m eps \u001b[39m=\u001b[39m nn_model(x_i, t_is, c_i)\n\u001b[1;32m     76\u001b[0m eps1 \u001b[39m=\u001b[39m eps[:n_sample]\n\u001b[1;32m     77\u001b[0m eps2 \u001b[39m=\u001b[39m eps[n_sample:]\n",
+      "File \u001b[0;32m/usr/local/pace-apps/manual/packages/pytorch/1.12.0/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
+      "Cell \u001b[0;32mIn[122], line 230\u001b[0m, in \u001b[0;36mContextUnet.forward\u001b[0;34m(self, x, timesteps, y)\u001b[0m\n\u001b[1;32m    228\u001b[0m \u001b[39m# print(\"0,h.shape =\", h.shape)\u001b[39;00m\n\u001b[1;32m    229\u001b[0m \u001b[39mfor\u001b[39;00m module \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39minput_blocks:\n\u001b[0;32m--> 230\u001b[0m     h \u001b[39m=\u001b[39m module(h, emb)\n\u001b[1;32m    231\u001b[0m     hs\u001b[39m.\u001b[39mappend(h)\n\u001b[1;32m    232\u001b[0m     \u001b[39m# print(\"module encoder, h.shape =\", h.shape)\u001b[39;00m\n\u001b[1;32m    233\u001b[0m \u001b[39m# print(\"2,h.shape =\", h.shape)\u001b[39;00m\n",
+      "File \u001b[0;32m/usr/local/pace-apps/manual/packages/pytorch/1.12.0/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
+      "Cell \u001b[0;32mIn[117], line 5\u001b[0m, in \u001b[0;36mTimestepEmbedSequential.forward\u001b[0;34m(self, x, emb, encoder_out)\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[39mfor\u001b[39;00m layer \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m:\n\u001b[1;32m      4\u001b[0m     \u001b[39mif\u001b[39;00m \u001b[39misinstance\u001b[39m(layer, TimestepBlock):\n\u001b[0;32m----> 5\u001b[0m         x \u001b[39m=\u001b[39m layer(x, emb)\n\u001b[1;32m      6\u001b[0m     \u001b[39melif\u001b[39;00m \u001b[39misinstance\u001b[39m(layer, AttentionBlock):\n\u001b[1;32m      7\u001b[0m         x \u001b[39m=\u001b[39m layer(x, encoder_out)\n",
       "File \u001b[0;32m/usr/local/pace-apps/manual/packages/pytorch/1.12.0/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
-      "Cell \u001b[0;32mIn[95], line 230\u001b[0m, in \u001b[0;36mContextUnet.forward\u001b[0;34m(self, x, timesteps, y)\u001b[0m\n\u001b[1;32m    228\u001b[0m \u001b[39m# print(\"0,h.shape =\", h.shape)\u001b[39;00m\n\u001b[1;32m    229\u001b[0m \u001b[39mfor\u001b[39;00m module \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39minput_blocks:\n\u001b[0;32m--> 230\u001b[0m     h \u001b[39m=\u001b[39m module(h, emb)\n\u001b[1;32m    231\u001b[0m     hs\u001b[39m.\u001b[39mappend(h)\n\u001b[1;32m    232\u001b[0m     \u001b[39m# print(\"module encoder, h.shape =\", h.shape)\u001b[39;00m\n\u001b[1;32m    233\u001b[0m \u001b[39m# print(\"2,h.shape =\", h.shape)\u001b[39;00m\n",
+      "Cell \u001b[0;32mIn[118], line 72\u001b[0m, in \u001b[0;36mResBlock.forward\u001b[0;34m(self, x, emb)\u001b[0m\n\u001b[1;32m     70\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m     71\u001b[0m     h \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m emb_out\n\u001b[0;32m---> 72\u001b[0m     h \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mout_layers(h)\n\u001b[1;32m     73\u001b[0m \u001b[39m# print(\"ResBlock, torch.unique(h).shape =\", torch.unique(h).shape)\u001b[39;00m\n\u001b[1;32m     74\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mskip_connection(x) \u001b[39m+\u001b[39m h\n",
       "File \u001b[0;32m/usr/local/pace-apps/manual/packages/pytorch/1.12.0/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
-      "Cell \u001b[0;32mIn[90], line 9\u001b[0m, in \u001b[0;36mTimestepEmbedSequential.forward\u001b[0;34m(self, x, emb, encoder_out)\u001b[0m\n\u001b[1;32m      7\u001b[0m         x \u001b[39m=\u001b[39m layer(x, encoder_out)\n\u001b[1;32m      8\u001b[0m     \u001b[39melse\u001b[39;00m:\n\u001b[0;32m----> 9\u001b[0m         x \u001b[39m=\u001b[39m layer(x)\n\u001b[1;32m     10\u001b[0m \u001b[39mreturn\u001b[39;00m x\n",
+      "File \u001b[0;32m/usr/local/pace-apps/manual/packages/pytorch/1.12.0/lib/python3.9/site-packages/torch/nn/modules/container.py:139\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    137\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39minput\u001b[39m):\n\u001b[1;32m    138\u001b[0m     \u001b[39mfor\u001b[39;00m module \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m:\n\u001b[0;32m--> 139\u001b[0m         \u001b[39minput\u001b[39m \u001b[39m=\u001b[39m module(\u001b[39minput\u001b[39;49m)\n\u001b[1;32m    140\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39minput\u001b[39m\n",
       "File \u001b[0;32m/usr/local/pace-apps/manual/packages/pytorch/1.12.0/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
       "File \u001b[0;32m/usr/local/pace-apps/manual/packages/pytorch/1.12.0/lib/python3.9/site-packages/torch/nn/modules/conv.py:457\u001b[0m, in \u001b[0;36mConv2d.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    456\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39minput\u001b[39m: Tensor) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m Tensor:\n\u001b[0;32m--> 457\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_conv_forward(\u001b[39minput\u001b[39;49m, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mweight, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbias)\n",
       "File \u001b[0;32m/usr/local/pace-apps/manual/packages/pytorch/1.12.0/lib/python3.9/site-packages/torch/nn/modules/conv.py:453\u001b[0m, in \u001b[0;36mConv2d._conv_forward\u001b[0;34m(self, input, weight, bias)\u001b[0m\n\u001b[1;32m    449\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mpadding_mode \u001b[39m!=\u001b[39m \u001b[39m'\u001b[39m\u001b[39mzeros\u001b[39m\u001b[39m'\u001b[39m:\n\u001b[1;32m    450\u001b[0m     \u001b[39mreturn\u001b[39;00m F\u001b[39m.\u001b[39mconv2d(F\u001b[39m.\u001b[39mpad(\u001b[39minput\u001b[39m, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_reversed_padding_repeated_twice, mode\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mpadding_mode),\n\u001b[1;32m    451\u001b[0m                     weight, bias, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstride,\n\u001b[1;32m    452\u001b[0m                     _pair(\u001b[39m0\u001b[39m), \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdilation, \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mgroups)\n\u001b[0;32m--> 453\u001b[0m \u001b[39mreturn\u001b[39;00m F\u001b[39m.\u001b[39;49mconv2d(\u001b[39minput\u001b[39;49m, weight, bias, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mstride,\n\u001b[1;32m    454\u001b[0m                 \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mpadding, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mdilation, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mgroups)\n",
-      "\u001b[0;31mRuntimeError\u001b[0m: Expected 3D (unbatched) or 4D (batched) input to conv2d, but got input of size: [2, 512]"
+      "\u001b[0;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 320.00 MiB (GPU 0; 23.64 GiB total capacity; 21.71 GiB already allocated; 170.50 MiB free; 22.45 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF"
      ]
     }
    ],