Update app.py

app.py

@@ -9,35 +9,32 @@ from torchvision import transforms
 # Make sure this matches the architecture you used to train your model
 
 
-from torchvision.models import resnet50
-import torch.nn as nn
-
-class MyCarClassifier(nn.Module):
-    def __init__(self, num_classes=196):
-        super(MyCarClassifier, self).__init__()
-        self.backbone = resnet50(weights=None)
-        in_ch = self.backbone.fc.in_features
-        self.backbone.fc = nn.Identity()
-        self.fc = nn.Sequential(
-            nn.Linear(in_ch, 1024),      # fc.1
-            nn.ReLU(),
-            nn.BatchNorm1d(1024),        # fc.3
-            nn.Linear(1024, 512),        # fc.5
-            nn.ReLU(),
-            nn.BatchNorm1d(512),         # fc.7
-            nn.Linear(512, num_classes)  # fc.9
-        )
-
-    def forward(self, x):
-        x = self.backbone(x)
-        x = self.fc(x)
-        return x
+# Model definition is now handled directly in the loading section
 
 # --------------------------
 # 2. Load model weights
 # --------------------------
-model = MyCarClassifier()
-model.load_state_dict(torch.load("best_stanford_cars_transfer_model.pth", map_location="cpu"), strict=True)
+# Load the checkpoint directly as it was saved (a plain ResNet50 with custom fc head)
+state_dict = torch.load("best_stanford_cars_transfer_model.pth", map_location="cpu")
+
+# Create a ResNet50 and modify its fc to match the checkpoint
+from torchvision.models import resnet50
+model = resnet50(weights=None)
+
+# Replace fc with the multi-layer head that matches checkpoint
+in_ch = model.fc.in_features
+model.fc = torch.nn.Sequential(
+    torch.nn.Linear(in_ch, 1024),      # fc.1
+    torch.nn.ReLU(),
+    torch.nn.BatchNorm1d(1024),        # fc.3  
+    torch.nn.Linear(1024, 512),        # fc.5
+    torch.nn.ReLU(),
+    torch.nn.BatchNorm1d(512),         # fc.7
+    torch.nn.Linear(512, 196)          # fc.9
+)
+
+# Load the state dict
+model.load_state_dict(state_dict, strict=True)
 model.eval()  # important for inference
 
 # --------------------------