update PR curve API usage

sematic-ai · Dec 12, 2022 · b20f52c · b20f52c
1 parent bea626e
commit b20f52c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sematic/examples/mnist/pytorch/train_eval.py b/sematic/examples/mnist/pytorch/train_eval.py
@@ -109,7 +109,7 @@ def test(model: nn.Module, device: torch.device, test_loader: DataLoader):
             correct += pred.eq(target).sum().item()
 
     test_loss /= len(test_loader.dataset)  # type: ignore
-    pr_curve = PrecisionRecallCurve(num_classes=10)
+    pr_curve = PrecisionRecallCurve(num_classes=10, task="multiclass")
     precision, recall, thresholds = pr_curve(torch.cat(probas), torch.cat(targets))
     classes = []
     for i in range(10):