Update app.py

app.py

@@ -56,8 +56,8 @@ class CNN(nn.Module):
         super(CNN, self).__init__()
         self.conv = nn.Conv2d(input_channels, 16, kernel_size=3, stride=1, padding=1)
         self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
-        # Adjusting the fully connected layer to accommodate the correct input size
-        self.fc = nn.Linear(16 * 4 * 4, output_dim)  # Adjust this based on your image size
+        # Adjust the fully connected layer to accommodate the correct input size
+        self.fc = nn.Linear(16 * 4 * 8, output_dim)  # 16 * 4 * 8 = 512
 
     def forward(self, x):
         x = self.pool(torch.relu(self.conv(x)))
@@ -65,6 +65,7 @@ class CNN(nn.Module):
         x = x.view(x.size(0), -1)  # Flatten for the fully connected layer
         return torch.sigmoid(self.fc(x))
 
+
 class PhiModel(nn.Module):
     def __init__(self, input_dim):
         super(PhiModel, self).__init__()