Update streamlit_app.py

streamlit_app.py

@@ -61,8 +61,14 @@ def load_model():
         model_state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
         num_classes = len(class_names)
 
-        # Detect model from classifier shape
-        classifier_input = model_state_dict['classifier.weight'].shape[1]
+        # Try to detect classifier shape; fallback if key missing
+        classifier_key = [k for k in model_state_dict.keys() if 'classifier.weight' in k or 'head.fc.weight' in k]
+        if not classifier_key:
+            st.error("Could not detect classifier layer in checkpoint.")
+            return None
+
+        classifier_input = model_state_dict[classifier_key[0]].shape[1]
+
         feature_map = {
             1280: 'efficientnet_b0',
             1408: 'efficientnet_b1',