Update app.py

app.py

@@ -82,12 +82,22 @@ class ModelManager:
                 num_classes=len(dog_breeds),
                 device=self.device
             ).to(self.device)
-
+    
             checkpoint = torch.load(
                 'ConvNextV2Base_best_model.pth',
                 map_location=self.device
             )
-            self._breed_model.load_state_dict(checkpoint['model_state_dict'], strict=False)
+            
+            # Try to load with model_state_dict first, then base_model
+            if 'model_state_dict' in checkpoint:
+                self._breed_model.load_state_dict(checkpoint['model_state_dict'], strict=False)
+            elif 'base_model' in checkpoint:
+                self._breed_model.load_state_dict(checkpoint['base_model'], strict=False)
+            else:
+                # If neither key exists, raise a descriptive error
+                available_keys = list(checkpoint.keys()) if isinstance(checkpoint, dict) else "not a dictionary"
+                raise KeyError(f"Model checkpoint does not contain 'model_state_dict' or 'base_model' keys. Available keys: {available_keys}")
+                
             self._breed_model.eval()
         return self._breed_model