Update inference.py

inference.py

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