proper mean calculation

Lightning-AI · Jan 3, 2025 · b67a5b6 · b67a5b6
1 parent 998a2bc
commit b67a5b6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/torchmetrics/functional/segmentation/mean_iou.py b/src/torchmetrics/functional/segmentation/mean_iou.py
@@ -112,4 +112,4 @@ def mean_iou(
     score = _mean_iou_compute(intersection, union)
     valid_classes = union > 0
     score = score * valid_classes
-    return score if per_class else score[valid_classes].mean(dim=-1)
+    return score if per_class else score.sum(dim=-1) / valid_classes.sum(dim=-1)