Fix: JaccardIndex multi-label compute (#1125)

(cherry picked from commit 883254e)

CHANGELOG.md

@@ -38,6 +38,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 -
 
 
+### Fixed
+
+-
+
+
+- Fixed JaccardIndex multi-label compute ([#1125](https://github.com/Lightning-AI/metrics/pull/1125))
+
+
+
 ## [0.9.2] - 2022-06-29
 
 ### Fixed

torchmetrics/classification/jaccard.py

@@ -104,10 +104,25 @@ def __init__(
 
     def compute(self) -> Tensor:
         """Computes intersection over union (IoU)"""
-        return _jaccard_from_confmat(
-            self.confmat,
-            self.num_classes,
-            self.average,
-            self.ignore_index,
-            self.absent_score,
-        )
+
+        if self.multilabel:
+            return torch.stack(
+                [
+                    _jaccard_from_confmat(
+                        confmat,
+                        2,
+                        self.average,
+                        self.ignore_index,
+                        self.absent_score,
+                    )[1]
+                    for confmat in self.confmat
+                ]
+            )
+        else:
+            return _jaccard_from_confmat(
+                self.confmat,
+                self.num_classes,
+                self.average,
+                self.ignore_index,
+                self.absent_score,
+            )