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()
|