Update app.py

app.py

@@ -97,7 +97,15 @@ 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.fc = torch.nn.Linear(model.fc.in_features, 3)  # Assuming 3 output classes
+        # Print model architecture to identify the classifier layer
+    print(model)  # Inspect the model structure
+
+    # Update the classifier layer to match the number of labels
+    if hasattr(model, 'classifier'):
+        model.classifier = torch.nn.Linear(model.classifier.in_features, 3)  # Assuming 3 output classes
+    else:
+        # Access the linear layer differently if 'classifier' does not exist
+
 
     model.train()