added missing non-linearity

Taeung · Jan 5, 2024 · 0d2a5c0 · 0d2a5c0
1 parent ed40a2c
commit 0d2a5c0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mnist/main.py b/mnist/main.py
@@ -20,8 +20,9 @@ def __init__(self):
 
     def forward(self, x):
         x = self.conv1(x)
-        x = F.relu(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)