added changes to model

app.py

@@ -28,6 +28,8 @@ device = torch.device('cpu')
 
 # Load your trained ResNet-50 model (or any custom architecture)
 model = models.resnet50(pretrained=False)  # Load the ResNet-50 architecture
+model.fc = nn.Linear(model.fc.in_features, 1000) 
+
 model.load_state_dict(torch.load("model.pth", map_location=device))  # Load the trained weights (.pth)
 model.to(device)  # Move model to CPU (even if you have a GPU)
 
@@ -42,8 +44,7 @@ transform = transforms.Compose([
 ])
 
 # Define the labels for ImageNet (or your specific dataset labels)
-LABELS = ["class_1", "class_2", "class_3", "class_4", "class_5",  # Replace with your classes
-          "class_6", "class_7", "class_8", "class_9", "class_10"]
+LABELS = [f"class_{k}" for k in range(1,1001)]
 
 # Define the prediction function
 def predict(image):