Repository: Lightning-AI/torchmetrics

Issue:
__call__ using ClassWrapper does not update wrapped metric
## 🐛 Bug

Using __call__ with ClasswiseWrapper does not update the Metric within, such that subsequent calls to compute are only using the last batch input.

### To Reproduce

See code sample

#### Code sample

```py
import numpy as np
import torch
import torchmetrics

np.random.seed(42)

m1 = torchmetrics.Precision(num_classes=3, average=None)
m2 = torchmetrics.Precision(num_classes=3, average=None)
mcw1 = torchmetrics.ClasswiseWrapper(torchmetrics.Precision(num_classes=3, average=None))
mcw2 = torchmetrics.ClasswiseWrapper(torchmetrics.Precision(num_classes=3, average=None))

for i in range(2):
    print(f"Epoch: {i}")
    preds = torch.Tensor(np.random.uniform(0, 1, (100, 3)))
    targets = torch.Tensor(np.random.randint(0, 2, (100, 3))).int()
    m1(preds, targets)
    m2.update(preds, targets)
    mcw1(preds, targets)
    mcw2.update(preds, targets)
    print(f"Prec:             {m1.compute().numpy()}")
    print(f"Prec update:      {m2.compute().numpy()}")
    print(f"Prec wrap:        {np.asarray([v.numpy() for v in mcw1.compute().values()])}")
    print(f"Prec wrap update: {np.asarray([v.numpy() for v in mcw2.compute().values()])}")
```

Output is:

```
Epoch: 0
Prec:             [0.42857143 0.537037   0.509434  ]
Prec update:      [0.42857143 0.537037   0.509434  ]
Prec wrap:        [0.42857143 0.537037   0.509434  ]
Prec wrap update: [0.42857143 0.537037   0.509434  ]
Epoch: 1
Prec:             [0.46464646 0.5247525  0.48039216]
Prec update:      [0.46464646 0.5247525  0.48039216]
Prec wrap:        [0.5       0.5106383 0.4489796]
Prec wrap update: [0.46464646 0.5247525  0.48039216]
```

### Expected behavior

The values should be the same whether calling `update` or `__call__`

### Environment

- TorchMetrics version: 0.9.3
- Python & PyTorch Version (e.g., 1.0): 1.12.0
- Any other relevant information such as OS (e.g., Linux): Ubuntu 20.04



Discussion:
Hi! thanks for your contribution!, great first issue!
I see this is because `full_state_update` is set to `True` inside the `ClasswiseWrapper` class. Perhaps the value of `full_state_update` should be set to that of the wrapped metric?

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
@@ -1,4 +1,17 @@
-from typing import Any, Dict, List, Optional
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Callable, Dict, List, Optional
 
 from torch import Tensor
 
@@ -51,8 +64,6 @@ class ClasswiseWrapper(Metric):
         'recall_horse': tensor(0.), 'recall_fish': tensor(0.3333), 'recall_dog': tensor(0.4000)}
     """
 
-    full_state_update: Optional[bool] = True
-
     def __init__(self, metric: Metric, labels: Optional[List[str]] = None) -> None:
         super().__init__()
         if not isinstance(metric, Metric):
@@ -61,6 +72,7 @@ def __init__(self, metric: Metric, labels: Optional[List[str]] = None) -> None:
             raise ValueError(f"Expected argument `labels` to either be `None` or a list of strings but got {labels}")
         self.metric = metric
         self.labels = labels
+        self._update_count = 1
 
     def _convert(self, x: Tensor) -> Dict[str, Any]:
         name = self.metric.__class__.__name__.lower()
@@ -68,6 +80,9 @@ def _convert(self, x: Tensor) -> Dict[str, Any]:
             return {f"{name}_{i}": val for i, val in enumerate(x)}
         return {f"{name}_{lab}": val for lab, val in zip(self.labels, x)}
 
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+        return self._convert(self.metric(*args, **kwargs))
+
     def update(self, *args: Any, **kwargs: Any) -> None:
         self.metric.update(*args, **kwargs)
 
@@ -76,3 +91,11 @@ def compute(self) -> Dict[str, Tensor]:
 
     def reset(self) -> None:
         self.metric.reset()
+
+    def _wrap_update(self, update: Callable) -> Callable:
+        """Overwrite to do nothing."""
+        return update
+
+    def _wrap_compute(self, compute: Callable) -> Callable:
+        """Overwrite to do nothing."""
+        return compute


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
@@ -18,8 +18,8 @@ def test_output_no_labels():
     base = Accuracy(num_classes=3, average=None)
     metric = ClasswiseWrapper(Accuracy(num_classes=3, average=None))
     for _ in range(2):
-        preds = torch.randn(10, 3).softmax(dim=-1)
-        target = torch.randint(3, (10,))
+        preds = torch.randn(20, 3).softmax(dim=-1)
+        target = torch.randint(3, (20,))
         val = metric(preds, target)
         val_base = base(preds, target)
         assert isinstance(val, dict)
@@ -35,8 +35,8 @@ def test_output_with_labels():
     base = Accuracy(num_classes=3, average=None)
     metric = ClasswiseWrapper(Accuracy(num_classes=3, average=None), labels=labels)
     for _ in range(2):
-        preds = torch.randn(10, 3).softmax(dim=-1)
-        target = torch.randint(3, (10,))
+        preds = torch.randn(20, 3).softmax(dim=-1)
+        target = torch.randint(3, (20,))
         val = metric(preds, target)
         val_base = base(preds, target)
         assert isinstance(val, dict)
@@ -44,6 +44,13 @@ def test_output_with_labels():
         for i, lab in enumerate(labels):
             assert f"accuracy_{lab}" in val
             assert val[f"accuracy_{lab}"] == val_base[i]
+        val = metric.compute()
+        val_base = base.compute()
+        assert isinstance(val, dict)
+        assert len(val) == 3
+        for i, lab in enumerate(labels):
+            assert f"accuracy_{lab}" in val
+            assert val[f"accuracy_{lab}"] == val_base[i]
 
 
 @pytest.mark.parametrize("prefix", [None, "pre_"])