Update app.py

app.py

@@ -17,11 +17,10 @@ for param in model.parameters():
 
 # # Replace the classifier layer
 # # to customise it according to our output
-# model.classifier = nn.Sequential(
-#     nn.Linear(256 * 7 * 7, 1024),
-#     nn.Linear(1024, 256),
-#     nn.Linear(256, 2),
-# )
+model.classifier = nn.Sequential(
+    nn.Linear(256 * 7 * 7, 1024),
+    nn.Linear(1024, 256),
+    nn.Linear(256, 2))
 
 checkpoint = torch.load(
     "CatVsDogsModel.pth", map_location=torch.device("cpu")