Fix checkpoint loading to extract model_state_dict properly

app.py

@@ -18,7 +18,11 @@ try:
     
     num_classes = len(label_to_class)
     model = FineGrainedClassifier(num_classes=num_classes)
-    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+    
+    # Load checkpoint and extract model_state_dict
+    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+    state_dict = checkpoint.get('model_state_dict', checkpoint)
+    model.load_state_dict(state_dict)
     model.eval()
     
     # Load text tokenizer