Update inference.py

inference.py

@@ -22,9 +22,11 @@ def predict(model, tokenizer, prompt, option1, option2, device):
     encoded = tokenizer(inputs, padding=True, truncation=True, return_tensors="pt").to(device)
 
     with torch.no_grad():
-        logits = model(encoded["input_ids"])  # shape: [2, 1]
+        outputs = model(encoded["input_ids"])  # shape: [2, 512]
 
-    probs = torch.softmax(logits.squeeze(), dim=0)  # shape: [2]
+    # Classifier expects input shape [2, 512] → [2, 1], so we stack scores
+    logits = model.classifier(outputs).squeeze(-1)  # shape: [2]
+    probs = torch.softmax(logits, dim=0)
     best = torch.argmax(probs).item()
 
     return {