Update utils.py

ChakradharG · Apr 23, 2024 · c6a9a8c · c6a9a8c
1 parent 8a00e1a
commit c6a9a8c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/model/utils.py b/model/utils.py
@@ -66,7 +66,7 @@ def test(model, dataloader, device):
         labels = batch['labels']
 
         with torch.inference_mode():
-            logits = model(inputs, attention_mask).squeeze()
+            logits = model(inputs, attention_mask).squeeze(1)
 
         preds = torch.where(logits > 0.5, 1, 0)
         predictions.extend(preds.cpu().tolist())