fix AUC test (allenai#4795)

ulmewennberg · Nov 17, 2020 · 3cad5b4 · 3cad5b4
1 parent efde092
commit 3cad5b4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/training/metrics/auc_test.py b/tests/training/metrics/auc_test.py
@@ -100,6 +100,7 @@ def test_distributed_auc(self):
         labels = torch.randint(3, 5, (8,), dtype=torch.long)
         # We make sure that the positive label is always present.
         labels[0] = 4
+        labels[4] = 4
 
         false_positive_rates, true_positive_rates, _ = metrics.roc_curve(
             labels.cpu().numpy(), predictions.cpu().numpy(), pos_label=4
@@ -124,6 +125,7 @@ def test_distributed_auc_unequal_batches(self):
         labels = torch.randint(3, 5, (8,), dtype=torch.long)
         # We make sure that the positive label is always present.
         labels[0] = 4
+        labels[4] = 4
 
         false_positive_rates, true_positive_rates, _ = metrics.roc_curve(
             labels.cpu().numpy(), predictions.cpu().numpy(), pos_label=4