Update streamlit_app.py

streamlit_app.py

@@ -37,22 +37,30 @@ def load_model():
         st.error("Model file not found. Please upload butterfly_classifier.pth to your space.")
         return None
     
-    # Load the checkpoint first to check the actual number of classes
+    # Load the checkpoint
     checkpoint = torch.load(MODEL_PATH, map_location="cpu")
     
-    # Get the number of classes from the saved model weights
-    if 'classifier.weight' in checkpoint:
-        num_classes_in_model = checkpoint['classifier.weight'].shape[0]
-    elif 'fc.weight' in checkpoint:  # Alternative naming
-        num_classes_in_model = checkpoint['fc.weight'].shape[0]
+    # Extract the actual model state dict
+    if 'model_state_dict' in checkpoint:
+        model_state_dict = checkpoint['model_state_dict']
+    else:
+        # If it's just the state dict directly
+        model_state_dict = checkpoint
+    
+    # Get the number of classes from the model weights
+    if 'classifier.weight' in model_state_dict:
+        num_classes_in_model = model_state_dict['classifier.weight'].shape[0]
+    elif 'fc.weight' in model_state_dict:
+        num_classes_in_model = model_state_dict['fc.weight'].shape[0]
     else:
         # Fallback: assume it matches class_names
         num_classes_in_model = len(class_names)
     
-    
-    # Create model with the correct number of classes from the saved model
+    # Create model with the correct number of classes
     model = timm.create_model('efficientnet_b0', pretrained=False, num_classes=num_classes_in_model)
-    model.load_state_dict(checkpoint)
+    
+    # Load the model state dict (not the entire checkpoint)
+    model.load_state_dict(model_state_dict)
     model.eval()
     
     return model