fix

tasks/text.py

@@ -109,7 +109,7 @@ def bert_classifier(test_dataset: dict, model: str):
             test_input_ids = batch["input_ids"].to(device)
             test_attention_mask = batch["attention_mask"].to(device)
             outputs = model(test_input_ids, test_attention_mask)
-            p = torch.argmax(outputs, dim=1)
+            p = torch.argmax(outputs.logits, dim=1)
             predictions = np.append(predictions, p.cpu().numpy())
     
     print("Finished BERT model run")
@@ -149,7 +149,7 @@ def moe_classifier(test_dataset: dict, model: str):
             embeddings = embedding_outputs.last_hidden_state[:, 0, :]
 
             outputs = model(embeddings)
-            p = torch.argmax(outputs.logits, dim=1)
+            p = torch.argmax(outputs, dim=1)
             predictions = np.append(predictions, p.cpu().numpy())
 
     print("Finished running MoE Classifier")