Update app.py

app.py

@@ -14,12 +14,12 @@ for param in model.parameters():
 
 # # Add a avgpool here
 # avgpool = nn.AdaptiveAvgPool2d((7, 7))
-
+# 
 # # 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(36, 256),
+    # nn.Linear(1024, 256),
     nn.Linear(256, 2))
 
 checkpoint = torch.load(
@@ -39,7 +39,7 @@ transform = transforms.Compose(
 
 def predict(img):
     img = transform(img).to('cpu')
-    # img = transform(img).unsqueeze(0)
+    img = transform(img).unsqueeze(0)
     with torch.no_grad():
         out=  model(img)
     probability = torch.nn.functional.softmax(out[0],dim=0)