Update app.py

app.py

@@ -60,10 +60,13 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Update this path with your model path
 model_path = 'best_model.pth'
-num_classes = 10  # Assuming you have 10 classes, update based on your dataset
+num_classes = 10  # Update based on your dataset
 
 model = EfficientNetWithNovelty(num_classes)
-model.load_state_dict(torch.load(model_path, map_location=device))
+
+# Load checkpoint instead of raw state_dict
+checkpoint = torch.load(model_path, map_location=device)
+model.load_state_dict(checkpoint["model_state_dict"])  # Correct key for model weights
 model.to(device)
 model.eval()