Update app.py

app.py

@@ -8,15 +8,19 @@ class IcebergClassifier(nn.Module):
     def __init__(self):
         super().__init__()
         self.conv = nn.Sequential(
-            nn.Conv2d(2,16,3), nn.ReLU(), nn.MaxPool2d(2),
-            nn.Conv2d(16,32,3), nn.ReLU(), nn.MaxPool2d(2)
+            nn.Conv2d(2, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
+            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2)
         )
         self.fc = nn.Sequential(
-            nn.Linear(32*17*17,64), nn.ReLU(),
-            nn.Linear(64,1), nn.Sigmoid()
+            nn.Linear(64 * 9 * 9, 64), nn.ReLU(), 
+            nn.Dropout(0.5),
+            nn.Linear(64, 1), 
+            nn.Sigmoid()
         )
+
     def forward(self, x): 
-        return self.fc(self.conv(x).view(x.size(0),-1))
+        return self.fc(self.conv(x).view(x.size(0), -1))
 
 @st.cache_resource
 def load_model():