diff --git a/TD2 Deep Learning.ipynb b/TD2 Deep Learning.ipynb
index 184e0396037195e6cdb030eef5db076be0ca4722..88c2cb7de2b9c94d56a889ecd22b390827780408 100644
--- a/TD2 Deep Learning.ipynb	
+++ b/TD2 Deep Learning.ipynb	
@@ -1111,43 +1111,52 @@
     "\n",
     "class Net_1(nn.Module):\n",
     "    def __init__(self):\n",
-    "        super(Net, self).__init__()\n",
+    "        super(Net_1, self).__init__()\n",
     "\n",
     "        #Define the 3 convolutional layers\n",
-    "\n",
     "        #First layer : 3 input channels, 16 output channels, kernel size 3, padding 1\n",
     "        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)\n",
-    "        self.pool1 = nn.MaxPool2d(2, 2) #MaxPool with kernel size 2 at first layer output\n",
-    "\n",
     "        #Second layer : 16 input channels, 32 output channels, kernel size 3, padding 1\n",
     "        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1)\n",
-    "        self.pool2 = nn.MaxPool2d(2, 2) #MaxPool with kernel size 2 at second layer output\n",
-    "\n",
     "        #Third layer : 32 input channels, 64 output channels, kernel size 3, padding 1\n",
-    "        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)\n",
-    "        self.pool2 = nn.MaxPool2d(2, 2) #MaxPool with kernel size 2 at third layer output\n",
+    "        self.conv3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)\n",
     "\n",
+    "        #MaxPool with kernel size 2\n",
+    "        self.pool = nn.MaxPool2d(2, 2)\n",
     "\n",
     "        #Define the 3 fully connected layers\n",
+    "        #First layer : input of size 64, image dimension 4*4, output of size 512\n",
+    "        self.fc1 = nn.Linear(in_features=64 * 4 * 4,out_features=512)\n",
+    "        #Second layer : input of size 512, output of size 64\n",
+    "        self.fc2 = nn.Linear(512, 64)\n",
+    "        #Third layer : input of size 64, output of size 10\n",
+    "        self.fc3 = nn.Linear(64, 10)\n",
+    "\n",
+    "        #Dropout\n",
+    "        self.dropout = nn.Dropout(p=0.5)\n",
     "\n",
-    "        #First layer\n",
-    "        self.fc1 = nn.Linear(16 * 5 * 5, 120)\n",
-    "        self.fc2 = nn.Linear(120, 84)\n",
-    "        self.fc3 = nn.Linear(84, 10)\n",
     "\n",
     "    def forward(self, x):\n",
+    "\n",
+    "        #Through the 3 convolutional layers\n",
     "        x = self.pool(F.relu(self.conv1(x)))\n",
     "        x = self.pool(F.relu(self.conv2(x)))\n",
-    "        x = x.view(-1, 16 * 5 * 5)\n",
-    "        x = F.relu(self.fc1(x))\n",
-    "        x = F.relu(self.fc2(x))\n",
+    "        x = self.pool(F.relu(self.conv3(x)))\n",
+    "\n",
+    "        #Linearize the output\n",
+    "        x = x.view(-1, 64 * 4 * 4) \n",
+    "\n",
+    "        #Through the 3 fully connected layers\n",
+    "        x = self.dropout(F.relu(self.fc1(x)))\n",
+    "        x = self.dropout(F.relu(self.fc2(x)))\n",
     "        x = self.fc3(x)\n",
+    "\n",
     "        return x\n",
     "\n",
     "\n",
     "# create a complete CNN\n",
-    "model = Net()\n",
-    "print(model)\n",
+    "model_1 = Net_1()\n",
+    "print(model_1)\n",
     "# move tensors to GPU if CUDA is available\n",
     "if train_on_gpu:\n",
     "    model.cuda()"