Update app.py

app.py

@@ -179,9 +179,16 @@ discriminator.eval()
 
 def get_prediction(embeddings):
     with torch.no_grad():
+        # Forward pass through the discriminator to get the logits and probabilities
         last_rep, logits, probs = discriminator(embeddings)
+        
+        # Filter logits to ignore the last dimension (assuming you only care about the first two)
+        filtered_logits = logits[:, 0:-1]
+
+        # Get the predicted labels using the filtered logits
+        _, predicted_labels = torch.max(filtered_logits, dim=-1)
 
-    predicted_labels = torch.argmax(probs,dim=-1)
+    # Convert to numpy array if needed
     predicted_labels = predicted_labels.cpu().numpy()
     return predicted_labels