Fix integration between classwise wrapper and metric tracker (#3004)

* fix integration between wrapper
* add tests
* changelog
* fix typing

(cherry picked from commit 0423a2c)

Author: SkafteNicki

CHANGELOG.md

@@ -21,7 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
--
+- Fixed integration between classwise wrapper and metric tracker ([#3004](https://github.com/PyTorchLightning/metrics/pull/3004))
 
 
 ## [1.6.2] - 2024-02-28

src/torchmetrics/wrappers/classwise.py

@@ -139,6 +139,11 @@ def __init__(
 
         self._update_count = 1
 
+    @property
+    def higher_is_better(self) -> Optional[bool]:  # type: ignore
+        """Return if the metric is higher the better."""
+        return self.metric.higher_is_better
+
     def _filter_kwargs(self, **kwargs: Any) -> dict[str, Any]:
         """Filter kwargs for the metric."""
         return self.metric._filter_kwargs(**kwargs)

src/torchmetrics/wrappers/tracker.py

@@ -24,6 +24,7 @@
 from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE
 from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, plot_single_or_multi_val
 from torchmetrics.utilities.prints import rank_zero_warn
+from torchmetrics.wrappers import ClasswiseWrapper
 
 if not _MATPLOTLIB_AVAILABLE:
     __doctest_skip__ = ["MetricTracker.plot"]
@@ -269,7 +270,7 @@ def best_metric(
                 return None, None
             return None
 
-        if isinstance(self._base_metric, Metric):
+        if isinstance(self._base_metric, Metric) and not isinstance(self._base_metric, ClasswiseWrapper):
             fn = torch.max if self.maximize else torch.min
             try:
                 value, idx = fn(res, 0)

tests/unittests/wrappers/test_tracker.py

@@ -26,7 +26,7 @@
 )
 from torchmetrics.regression import MeanAbsoluteError, MeanSquaredError
 from torchmetrics.utilities.imports import _TORCHMETRICS_GREATER_EQUAL_1_6
-from torchmetrics.wrappers import MetricTracker, MultioutputWrapper
+from torchmetrics.wrappers import ClasswiseWrapper, MetricTracker, MultioutputWrapper
 from unittests._helpers import seed_all
 
 seed_all(42)
@@ -101,6 +101,11 @@ def test_raises_error_if_increment_not_called(method, method_input):
             (torch.randn(50), torch.randn(50)),
             [False, False],
         ),
+        (
+            ClasswiseWrapper(MulticlassAccuracy(num_classes=3, average=None)),
+            (torch.randint(3, (50,)), torch.randint(3, (50,))),
+            True,
+        ),
     ],
 )
 def test_tracker(base_metric, metric_input, maximize):
@@ -244,6 +249,7 @@ def test_tracker_futurewarning():
         MeanAbsoluteError(),
         MulticlassAccuracy(num_classes=10),
         MetricCollection([MeanSquaredError(), MeanAbsoluteError()]),
+        ClasswiseWrapper(MulticlassAccuracy(num_classes=10, average=None)),
     ],
 )
 def test_tracker_higher_is_better_integration(base_metric):