Update app.py

app.py

@@ -16,7 +16,7 @@ class Net(nn.Module):
         self.fc3 = nn.Linear(64, 10)
 
     def forward(self, x):
-        x = x.view(-1, 28 * 28)
+        x = x.view(-1, 28 * 28)  # Flatten the input
         x = F.relu(self.fc1(x))
         x = F.relu(self.fc2(x))
         x = self.fc3(x)