Fix binary classification

openvinotoolkit · Sep 4, 2024 · ed18cbe · ed18cbe
1 parent 112b2b2
commit ed18cbe
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/otx/core/metrics/accuracy.py b/src/otx/core/metrics/accuracy.py
@@ -346,8 +346,10 @@ def compute(self) -> torch.Tensor:
 
 
 def _multi_class_cls_metric_callable(label_info: LabelInfo) -> MetricCollection:
+    num_classes = label_info.num_classes
+    task = "binary" if num_classes == 1 else "multiclass"
     return MetricCollection(
-        {"accuracy": TorchmetricAcc(task="multiclass", num_classes=label_info.num_classes)},
+        {"accuracy": TorchmetricAcc(task=task, num_classes=num_classes)},
     )