remove logits

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.logits, dim=1)
+            p = torch.argmax(outputs, dim=1)
             predictions = np.append(predictions, p.cpu().numpy())
     
     print("Finished BERT model run")