added missing non-linearity

ShirleneLiu · Mar 5, 2020 · dce9319 · dce9319
1 parent 47f09db
commit dce9319
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/mnist/main.py b/mnist/main.py
@@ -22,6 +22,7 @@ def forward(self, x):
         x = self.conv1(x)
         x = F.relu(x)
         x = self.conv2(x)
+        x = F.relu(x)
         x = F.max_pool2d(x, 2)
         x = self.dropout1(x)
         x = torch.flatten(x, 1)