Update app.py

app.py

@@ -11,7 +11,7 @@ def load_model(model_path, device):
     model = torchvision.models.densenet201(weights=weights).to(device)
     model.classifier = torch.nn.Sequential(
         torch.nn.Dropout(p=0.2, inplace=True),
-        torch.nn.Linear(in_features=1920, out_features=4, bias=True)
+        torch.nn.Linear(in_features=1920, out_features=2, bias=True)
     ).to(device)
     model.load_state_dict(torch.load(model_path, map_location=device))
     model.to(device)