Update app.py

app.py

@@ -102,11 +102,12 @@ def fine_tune_classification_model(train_loader):
 
     # 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
+        if isinstance(model.classifier, torch.nn.Sequential):
+            model.classifier[-1] = torch.nn.Linear(model.classifier[-1].in_features, 3)  # Assuming 3 output classes
+        else:
+            model.classifier = torch.nn.Linear(model.classifier.in_features, 3)  # In case it's a Linear layer directly
     else:
-        # Access the linear layer differently if 'classifier' does not exist
-        model.classifier = torch.nn.Linear(model.config.num_labels, 3)  # Update according to available layers
-
+        print("Classifier layer not found")
 
     model.train()