Update app.py
Browse files
app.py
CHANGED
|
@@ -11,9 +11,8 @@ transform = transforms.Compose([
|
|
| 11 |
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
| 12 |
])
|
| 13 |
|
| 14 |
-
model = models.
|
| 15 |
-
|
| 16 |
-
model.fc = nn.Linear(num_ftrs, 2)
|
| 17 |
model = model.to("cpu")
|
| 18 |
model.load_state_dict(torch.load("cnn_model.pth", weights_only=True, map_location="cpu"))
|
| 19 |
model.eval()
|
|
|
|
| 11 |
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
| 12 |
])
|
| 13 |
|
| 14 |
+
model = models.mobilenet_v3_large(pretrained=True)
|
| 15 |
+
model.classifier[3] = nn.Linear(model.classifier[3].in_features, 2)
|
|
|
|
| 16 |
model = model.to("cpu")
|
| 17 |
model.load_state_dict(torch.load("cnn_model.pth", weights_only=True, map_location="cpu"))
|
| 18 |
model.eval()
|