Repository: Lightning-AI/torchmetrics

Issue:
Add `_filter_kwargs` in `ClasswiseWrapper` metric wrapper
As the `ClasswiseWrapper` accepts agnostic `*args, **kwargs` arguments, its not handy to use it in a `MetricColletion`, as tailored input arguments for other metrics get passed in the metrics.

It would be great, just like in `MetricColletion`, to add the `_filter_kwargs` for the metric 


Discussion:
Hi! thanks for your contribution!, great first issue!
> It would be great, just like in `MetricColletion`, to add the `_filter_kwargs` for the metric

sounds fine to me... what do you think @justusschock @SkafteNicki ?

Resolving patch:
diff --git a/src/torchmetrics/wrappers/classwise.py b/src/torchmetrics/wrappers/classwise.py
--- a/src/torchmetrics/wrappers/classwise.py
+++ b/src/torchmetrics/wrappers/classwise.py
@@ -142,7 +142,12 @@ def __init__(
 
         self._update_count = 1
 
-    def _convert(self, x: Tensor) -> Dict[str, Any]:
+    def _filter_kwargs(self, **kwargs: Any) -> Dict[str, Any]:
+        """Filter kwargs for the metric."""
+        return self.metric._filter_kwargs(**kwargs)
+
+    def _convert_output(self, x: Tensor) -> Dict[str, Any]:
+        """Convert output to dictionary with labels as keys."""
         # Will set the class name as prefix if neither prefix nor postfix is given
         if not self._prefix and not self._postfix:
             prefix = f"{self.metric.__class__.__name__.lower()}_"
@@ -156,7 +161,7 @@ def _convert(self, x: Tensor) -> Dict[str, Any]:
 
     def forward(self, *args: Any, **kwargs: Any) -> Any:
         """Calculate on batch and accumulate to global state."""
-        return self._convert(self.metric(*args, **kwargs))
+        return self._convert_output(self.metric(*args, **kwargs))
 
     def update(self, *args: Any, **kwargs: Any) -> None:
         """Update state."""
@@ -164,7 +169,7 @@ def update(self, *args: Any, **kwargs: Any) -> None:
 
     def compute(self) -> Dict[str, Tensor]:
         """Compute metric."""
-        return self._convert(self.metric.compute())
+        return self._convert_output(self.metric.compute())
 
     def reset(self) -> None:
         """Reset metric."""


Test patch:
diff --git a/tests/unittests/wrappers/test_classwise.py b/tests/unittests/wrappers/test_classwise.py
--- a/tests/unittests/wrappers/test_classwise.py
+++ b/tests/unittests/wrappers/test_classwise.py
@@ -2,6 +2,7 @@
 import torch
 from torchmetrics import MetricCollection
 from torchmetrics.classification import MulticlassAccuracy, MulticlassF1Score, MulticlassRecall
+from torchmetrics.clustering import CalinskiHarabaszScore
 from torchmetrics.wrappers import ClasswiseWrapper
 
 
@@ -150,3 +151,27 @@ def test_double_use_of_prefix_with_metriccollection():
     assert "val/accuracy" in res
     assert "val/f_score_Tree" in res
     assert "val/f_score_Bush" in res
+
+
+def test_filter_kwargs_and_metriccollection():
+    """Test that kwargs are correctly filtered when using metric collection."""
+    metric = MetricCollection(
+        {
+            "accuracy": ClasswiseWrapper(MulticlassAccuracy(num_classes=3, average=None)),
+            "cluster": CalinskiHarabaszScore(),
+        },
+    )
+    preds = torch.randn(10, 3).softmax(dim=-1)
+    target = torch.randint(3, (10,))
+    data = torch.randn(10, 3)
+
+    metric.update(preds=preds, target=target, data=data, labels=target)
+    metric(preds=preds, target=target, data=data, labels=target)
+    val = metric.compute()
+
+    assert isinstance(val, dict)
+    assert len(val) == 4
+    assert "multiclassaccuracy_0" in val
+    assert "multiclassaccuracy_1" in val
+    assert "multiclassaccuracy_2" in val
+    assert "cluster" in val