Update app.py

app.py

@@ -105,6 +105,7 @@ def fine_tune_classification_model(train_loader):
         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.classifier = torch.nn.Linear(model.config.num_labels, 3)  # Update according to available layers
 
 
     model.train()