Update app.py

app.py

@@ -97,7 +97,7 @@ class CustomImageDataset(Dataset):
 def fine_tune_classification_model(train_loader):
     # Load the ResNet model with ignore_mismatched_sizes
     model = ResNetForImageClassification.from_pretrained('microsoft/resnet-50', num_labels=3, ignore_mismatched_sizes=True)
-    model.classifier = torch.nn.Linear(model.config.hidden_size, 3)  # Update classifier for 3 labels
+    model.classifier = torch.nn.Linear(model.config.num_features, 3)  # Update classifier for 3 labels
     model.train()
     
     optimizer = AdamW(model.parameters(), lr=1e-4)