From 57cf63427067e5369128ae849b28102a41549ec5 Mon Sep 17 00:00:00 2001
From: Miguel <138483675+mgp87@users.noreply.github.com>
Date: Thu, 1 Feb 2024 09:29:22 +0100
Subject: [PATCH] Add files via upload

---
 Unet.ipynb | 313 +++++++++++++++++++++++++++++++++++++++--------------
 unet.py    |  10 +-
 2 files changed, 237 insertions(+), 86 deletions(-)

diff --git a/Unet.ipynb b/Unet.ipynb
index f5564bc..6ddc23a 100644
--- a/Unet.ipynb
+++ b/Unet.ipynb
@@ -18,7 +18,7 @@
   "cells": [
     {
       "cell_type": "code",
-      "execution_count": 1,
+      "execution_count": 7,
       "metadata": {
         "id": "FJ8kS6a59r7M"
       },
@@ -37,10 +37,10 @@
         "  def __init__(self, in_channels, out_channels):\n",
         "    super(DoubleConv, self).__init__()\n",
         "    self.conv = nn.Sequential(\n",
-        "        nn.Conv2d(in_channels, out_channels, 3, 1),\n",
+        "        nn.Conv2d(in_channels, out_channels, 3, 1, 1),\n",
         "        nn.BatchNorm2d(out_channels),\n",
         "        nn.ReLU(inplace=True),\n",
-        "        nn.Conv2d(out_channels, out_channels, 3, 1),\n",
+        "        nn.Conv2d(out_channels, out_channels, 3, 1, 1),\n",
         "        nn.BatchNorm2d(out_channels),\n",
         "        nn.ReLU(inplace=True)\n",
         "    )\n",
@@ -51,7 +51,7 @@
       "metadata": {
         "id": "Te6UN_al-A7i"
       },
-      "execution_count": 2,
+      "execution_count": 8,
       "outputs": []
     },
     {
@@ -66,13 +66,13 @@
         "\n",
         "    # Down the U network\n",
         "    for feature in features:\n",
-        "      self.down_sampling.append(DoubleConv(in_channels, out_channels))\n",
+        "      self.down_sampling.append(DoubleConv(in_channels, feature))\n",
         "      in_channels = feature\n",
         "\n",
         "    # Up the U network\n",
         "    for feature in reversed(features):\n",
         "      self.up_sampling.append(nn.ConvTranspose2d(feature * 2, feature, kernel_size=2, stride=2))\n",
-        "      self.up_sampling.append(DoubleConv(feature * 2, out_channels))\n",
+        "      self.up_sampling.append(DoubleConv(feature * 2, feature))\n",
         "\n",
         "    self.bottle_neck = DoubleConv(features[-1], features[-1] * 2)\n",
         "    self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)\n",
@@ -91,7 +91,7 @@
         "      x = self.up_sampling[i](x)\n",
         "      skip_connection = skip_connections[i//2]\n",
         "\n",
-        "      if x.shape != skip_connection.shape():\n",
+        "      if x.shape != skip_connection.shape:\n",
         "        x = TF.resize(x, size=skip_connection.shape[2:])\n",
         "\n",
         "      concat_skip = torch.cat((skip_connection, x), dim=1)\n",
@@ -102,7 +102,7 @@
       "metadata": {
         "id": "E_e_pf2M_jlo"
       },
-      "execution_count": 3,
+      "execution_count": 15,
       "outputs": []
     },
     {
@@ -117,39 +117,39 @@
           "base_uri": "https://localhost:8080/"
         },
         "id": "WW-hRFUsPfrz",
-        "outputId": "2be2be3d-0175-4501-8cb2-0cd783b859a8"
+        "outputId": "b0cd763c-e1a0-4c79-95c1-0c9b70155147"
       },
-      "execution_count": 4,
+      "execution_count": 16,
       "outputs": [
         {
           "output_type": "stream",
           "name": "stdout",
           "text": [
-            "tensor([[[[ 0.4745,  0.3990,  0.4629,  ..., -0.4358, -2.7471,  1.2854],\n",
-            "          [-0.8002, -1.3700,  0.1394,  ..., -1.0008, -0.1001,  0.3647],\n",
-            "          [ 0.4885, -0.2957,  0.3457,  ..., -0.8069,  0.0458, -1.3301],\n",
+            "tensor([[[[-1.0468, -0.0166,  0.0301,  ...,  1.6915,  0.5362, -1.6300],\n",
+            "          [ 0.9801, -0.3030,  0.6119,  ...,  0.9953,  0.8321, -0.6704],\n",
+            "          [-0.0898,  0.1024,  2.3816,  ...,  0.0412, -0.6030, -0.6178],\n",
             "          ...,\n",
-            "          [-0.5858, -0.7971,  0.6741,  ...,  1.6228, -0.3367,  2.7377],\n",
-            "          [-0.6080, -1.0116,  0.5573,  ...,  0.4941,  1.8874,  1.1634],\n",
-            "          [-1.2275,  0.4125, -0.2102,  ...,  0.1442,  0.3863, -0.5045]]],\n",
+            "          [-0.0446, -0.5421, -0.9536,  ...,  2.4646, -0.6498, -0.7474],\n",
+            "          [-0.2652,  0.1830, -2.3945,  ...,  0.1613, -1.1950, -1.2325],\n",
+            "          [ 0.6813, -0.5447, -0.8322,  ...,  0.2796, -0.8650, -1.0577]]],\n",
             "\n",
             "\n",
-            "        [[[ 0.5432,  1.7162, -1.2448,  ..., -1.7358,  0.5377,  0.3850],\n",
-            "          [ 0.2976, -0.4702,  0.3024,  ..., -0.3201,  1.1905, -0.1385],\n",
-            "          [ 0.3186, -0.6002, -0.4442,  ...,  2.2087, -0.7032,  0.0036],\n",
+            "        [[[ 1.0751,  0.0246, -0.5136,  ...,  0.4104,  0.2865, -0.2959],\n",
+            "          [ 0.7612, -0.4717,  0.6451,  ...,  0.3040, -0.0145, -0.9153],\n",
+            "          [-1.0468, -0.4714,  0.6018,  ...,  0.0491,  1.7783,  0.4017],\n",
             "          ...,\n",
-            "          [-2.1623,  0.3663, -0.7531,  ...,  0.3838, -0.3361, -1.0889],\n",
-            "          [ 0.4942, -0.4374, -0.7796,  ..., -1.2864,  1.8577, -0.3994],\n",
-            "          [-0.2362, -1.0003, -0.6744,  ...,  1.7848,  0.4088, -0.0438]]],\n",
+            "          [ 1.5515, -2.0278, -1.8483,  ...,  0.4439,  0.5086, -0.5200],\n",
+            "          [-0.5622,  0.4031, -0.9494,  ..., -1.5475,  2.1311, -1.2994],\n",
+            "          [ 1.0823,  0.2312,  0.7589,  ..., -0.4231,  1.1247, -0.7772]]],\n",
             "\n",
             "\n",
-            "        [[[-0.1844,  0.4704, -1.1642,  ...,  0.0059,  0.5202, -1.2484],\n",
-            "          [ 0.6056,  0.0775, -0.6084,  ..., -0.3822,  1.0274,  0.7384],\n",
-            "          [-2.2475,  1.0625,  0.9412,  ..., -0.8918,  1.0682,  0.2734],\n",
+            "        [[[ 1.3303,  0.5727, -0.0660,  ...,  0.2148, -1.4475, -1.1001],\n",
+            "          [-0.7953, -1.0506,  1.0483,  ..., -0.0606, -0.4288, -1.0632],\n",
+            "          [-0.6977,  0.6204,  1.2682,  ..., -0.5036,  0.3979,  0.8696],\n",
             "          ...,\n",
-            "          [-1.4826,  0.1044,  0.0215,  ..., -1.0512,  0.2867, -0.1331],\n",
-            "          [-0.2543, -1.4832, -2.3471,  ...,  0.0743,  0.5073,  0.2504],\n",
-            "          [ 0.4399, -1.3677, -0.2083,  ...,  1.6405, -1.2264,  0.1492]]]])\n"
+            "          [ 1.2259, -0.9382,  0.6458,  ...,  0.0640, -2.0215,  0.6428],\n",
+            "          [ 0.1432,  0.1387,  0.3388,  ..., -0.6358, -1.2691,  0.2593],\n",
+            "          [ 0.4122,  0.2079,  1.0262,  ..., -0.2089, -0.8159,  0.4246]]]])\n"
           ]
         }
       ]
@@ -163,36 +163,68 @@
       ],
       "metadata": {
         "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 363
+          "base_uri": "https://localhost:8080/"
         },
         "id": "kwkddAv2OADq",
-        "outputId": "e468eb29-8786-40f1-fd8f-e0cd4e9a1e0c"
+        "outputId": "779aef22-7a41-4e2b-9ded-2e01a6ab6cd7"
       },
-      "execution_count": 5,
+      "execution_count": 17,
       "outputs": [
         {
-          "output_type": "error",
-          "ename": "RuntimeError",
-          "evalue": "Given groups=1, weight of size [1, 64, 3, 3], expected input[3, 1, 80, 80] to have 64 channels, but got 1 channels instead",
-          "traceback": [
-            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
-            "\u001b[0;32m<ipython-input-5-8e0f536ae55a>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mpreds\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpreds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m<ipython-input-3-35d19587483a>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     22\u001b[0m     \u001b[0mskip_connections\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mdown_sample\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdown_sampling\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 24\u001b[0;31m       \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdown_sample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     25\u001b[0m       \u001b[0mskip_connections\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m       \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax_pool\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m<ipython-input-2-85492b3c2b57>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    213\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    214\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 215\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    216\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    217\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    458\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    459\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 460\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_conv_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    461\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    462\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mConv3d\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_ConvNd\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py\u001b[0m in \u001b[0;36m_conv_forward\u001b[0;34m(self, input, weight, bias)\u001b[0m\n\u001b[1;32m    454\u001b[0m                             \u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstride\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    455\u001b[0m                             _pair(0), self.dilation, self.groups)\n\u001b[0;32m--> 456\u001b[0;31m         return F.conv2d(input, weight, bias, self.stride,\n\u001b[0m\u001b[1;32m    457\u001b[0m                         self.padding, self.dilation, self.groups)\n\u001b[1;32m    458\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;31mRuntimeError\u001b[0m: Given groups=1, weight of size [1, 64, 3, 3], expected input[3, 1, 80, 80] to have 64 channels, but got 1 channels instead"
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).\n",
+            "  warnings.warn(\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "tensor([[[[-7.8615e-01, -9.2691e-01, -6.0194e-01,  ..., -6.1102e-01,\n",
+            "           -5.9466e-01, -4.2123e-01],\n",
+            "          [-3.3110e-01,  5.0412e-02, -8.5804e-02,  ..., -3.5462e-01,\n",
+            "           -4.5124e-01, -2.9182e-01],\n",
+            "          [-6.5551e-02, -4.8520e-01, -1.3581e+00,  ..., -1.8820e-01,\n",
+            "           -1.5520e-01,  3.1639e-03],\n",
+            "          ...,\n",
+            "          [-5.7518e-01,  1.7181e-01, -9.2582e-02,  ..., -6.8580e-02,\n",
+            "           -1.9454e-01, -4.0080e-01],\n",
+            "          [-9.4960e-01, -3.5209e-01, -3.9878e-01,  ..., -3.0820e-01,\n",
+            "            1.5029e-01, -3.1823e-01],\n",
+            "          [-7.8362e-01, -9.5320e-01, -3.1767e-01,  ..., -3.7102e-01,\n",
+            "           -6.0932e-01, -8.3625e-01]]],\n",
+            "\n",
+            "\n",
+            "        [[[-5.1371e-01, -7.1526e-01, -5.1797e-01,  ..., -9.2383e-01,\n",
+            "           -8.4946e-01, -3.7277e-01],\n",
+            "          [-8.2932e-01, -1.7809e-01, -3.8386e-01,  ..., -6.7520e-01,\n",
+            "           -5.2506e-01, -2.0573e-01],\n",
+            "          [-9.3634e-02, -9.0712e-01, -3.1235e-01,  ..., -5.0498e-01,\n",
+            "           -5.9980e-01, -4.7795e-01],\n",
+            "          ...,\n",
+            "          [-5.8407e-01,  2.1302e-01,  7.1952e-02,  ...,  4.0669e-01,\n",
+            "           -4.8185e-01, -2.7956e-01],\n",
+            "          [-7.4300e-01, -9.4989e-01, -4.8291e-01,  ..., -1.4906e-02,\n",
+            "           -1.0740e-01, -1.5927e-01],\n",
+            "          [-5.2224e-01, -5.9556e-01, -2.2506e-01,  ..., -1.2161e+00,\n",
+            "           -5.8694e-01, -4.4059e-01]]],\n",
+            "\n",
+            "\n",
+            "        [[[-5.9068e-01, -1.0337e+00, -4.0911e-01,  ..., -5.8985e-01,\n",
+            "           -7.9712e-01, -3.6009e-01],\n",
+            "          [-3.1342e-04, -2.4367e-01, -6.4197e-02,  ..., -2.6382e-01,\n",
+            "           -3.8880e-01, -1.9363e-01],\n",
+            "          [-5.0104e-01, -2.7902e-01, -4.5525e-01,  ..., -5.6507e-01,\n",
+            "           -6.4159e-01, -1.2957e-01],\n",
+            "          ...,\n",
+            "          [-5.1757e-01, -3.9730e-01, -1.6993e-01,  ..., -1.7030e-01,\n",
+            "            6.7767e-02,  8.3066e-02],\n",
+            "          [-4.7388e-01, -8.4787e-01, -4.0920e-01,  ..., -4.6591e-01,\n",
+            "           -2.0471e-01, -2.8354e-01],\n",
+            "          [-4.0559e-01, -5.3141e-01, -5.2516e-01,  ..., -4.3604e-02,\n",
+            "           -1.0437e+00, -8.7411e-01]]]], grad_fn=<ConvolutionBackward0>)\n"
           ]
         }
       ]
@@ -209,7 +241,7 @@
       "metadata": {
         "id": "lZb0-YZdRz9c"
       },
-      "execution_count": 6,
+      "execution_count": 18,
       "outputs": []
     },
     {
@@ -233,7 +265,7 @@
       "metadata": {
         "id": "564AC2r1SN90"
       },
-      "execution_count": 7,
+      "execution_count": 19,
       "outputs": []
     },
     {
@@ -246,7 +278,7 @@
       "metadata": {
         "id": "TZC5dM_BTCWB"
       },
-      "execution_count": 8,
+      "execution_count": 20,
       "outputs": []
     },
     {
@@ -267,36 +299,155 @@
       ],
       "metadata": {
         "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 1000
+          "base_uri": "https://localhost:8080/"
         },
         "id": "DReFEqITTRFR",
-        "outputId": "7b37bf10-52a8-499c-82e7-6afcaa24b47e"
+        "outputId": "aba74d6a-0ed9-41da-ca40-1757b4faebd2"
       },
-      "execution_count": 9,
+      "execution_count": null,
       "outputs": [
         {
-          "output_type": "error",
-          "ename": "RuntimeError",
-          "evalue": "Given groups=1, weight of size [1, 64, 3, 3], expected input[6, 1, 126, 126] to have 64 channels, but got 1 channels instead",
-          "traceback": [
-            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
-            "\u001b[0;32m<ipython-input-9-0fab0316e20b>\u001b[0m in \u001b[0;36m<cell line: 4>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      6\u001b[0m     \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m     \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0munet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      9\u001b[0m     \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriteria\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtargets\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m     \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m<ipython-input-3-35d19587483a>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     22\u001b[0m     \u001b[0mskip_connections\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mdown_sample\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdown_sampling\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 24\u001b[0;31m       \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdown_sample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     25\u001b[0m       \u001b[0mskip_connections\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m       \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax_pool\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m<ipython-input-2-85492b3c2b57>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    213\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    214\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 215\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    216\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    217\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    458\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    459\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 460\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_conv_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbias\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    461\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    462\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mConv3d\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_ConvNd\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/conv.py\u001b[0m in \u001b[0;36m_conv_forward\u001b[0;34m(self, input, weight, bias)\u001b[0m\n\u001b[1;32m    454\u001b[0m                             \u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstride\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    455\u001b[0m                             _pair(0), self.dilation, self.groups)\n\u001b[0;32m--> 456\u001b[0;31m         return F.conv2d(input, weight, bias, self.stride,\n\u001b[0m\u001b[1;32m    457\u001b[0m                         self.padding, self.dilation, self.groups)\n\u001b[1;32m    458\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
-            "\u001b[0;31mRuntimeError\u001b[0m: Given groups=1, weight of size [1, 64, 3, 3], expected input[6, 1, 126, 126] to have 64 channels, but got 1 channels instead"
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n",
+            "Epoch: 0, Loss: -0.0\n"
           ]
         }
       ]
diff --git a/unet.py b/unet.py
index 5b5944e..aa8ee36 100644
--- a/unet.py
+++ b/unet.py
@@ -16,10 +16,10 @@ class DoubleConv(nn.Module):
   def __init__(self, in_channels, out_channels):
     super(DoubleConv, self).__init__()
     self.conv = nn.Sequential(
-        nn.Conv2d(in_channels, out_channels, 3, 1),
+        nn.Conv2d(in_channels, out_channels, 3, 1, 1),
         nn.BatchNorm2d(out_channels),
         nn.ReLU(inplace=True),
-        nn.Conv2d(out_channels, out_channels, 3, 1),
+        nn.Conv2d(out_channels, out_channels, 3, 1, 1),
         nn.BatchNorm2d(out_channels),
         nn.ReLU(inplace=True)
     )
@@ -36,13 +36,13 @@ def __init__(self, in_channels=3, out_channels=1, features=[64, 128, 256, 512]):
 
     # Down the U network
     for feature in features:
-      self.down_sampling.append(DoubleConv(in_channels, out_channels))
+      self.down_sampling.append(DoubleConv(in_channels, feature))
       in_channels = feature
 
     # Up the U network
     for feature in reversed(features):
       self.up_sampling.append(nn.ConvTranspose2d(feature * 2, feature, kernel_size=2, stride=2))
-      self.up_sampling.append(DoubleConv(feature * 2, out_channels))
+      self.up_sampling.append(DoubleConv(feature * 2, feature))
 
     self.bottle_neck = DoubleConv(features[-1], features[-1] * 2)
     self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
@@ -61,7 +61,7 @@ def forward(self, x):
       x = self.up_sampling[i](x)
       skip_connection = skip_connections[i//2]
 
-      if x.shape != skip_connection.shape():
+      if x.shape != skip_connection.shape:
         x = TF.resize(x, size=skip_connection.shape[2:])
 
       concat_skip = torch.cat((skip_connection, x), dim=1)