update

tasks/text.py

@@ -89,13 +89,12 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     with torch.no_grad():
         print('BEFORE PREDICTION')
-        for batch in dataloader:
-            print('INSIDE PREDICTION')
-            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)
-            predictions = np.append(predictions, p.cpu().numpy())
+        
+        test_input_ids = tokenized_test["input_ids"].to(device)
+        test_attention_mask = tokenized_test["attention_mask"].to(device)
+        outputs = model(test_input_ids, test_attention_mask)
+        p = torch.argmax(outputs.logits, dim=1)
+        predictions = np.append(predictions, p.cpu().numpy())
 
     print("Finished prediction run")