Update app.py

app.py

@@ -22,15 +22,19 @@ class ModifiedLargeNet(nn.Module):
     def __init__(self):
         super(ModifiedLargeNet, self).__init__()
         self.name = "modified_large"
-        self.fc1 = nn.Linear(128 * 128 * 3, 256)
-        self.fc2 = nn.Linear(256, 128)
-        self.fc3 = nn.Linear(128, 3)  # 3 classes: Rope, Hammer, Other
+        self.conv1 = nn.Conv2d(3, 5, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(5, 10, 5)
+        self.fc1 = nn.Linear(10 * 29 * 29, 32)
+        self.fc2 = nn.Linear(32, 3) # classify into "Rope"/"Hammer"/"others"
 
     def forward(self, x):
-        x = x.view(-1, 128 * 128 * 3)
-        x = torch.relu(self.fc1(x))
-        x = torch.relu(self.fc2(x))
-        x = self.fc3(x)
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 10 * 29 * 29)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        x = x.squeeze(1) # Flatten to [batch_size]
         return x