Repository: Lightning-AI/torchmetrics

Issue:
CompositionalMetric only returns value for last batch
## 🐛 Bug

When using a compositional metric, i.e., `metric_a + metric_b` or `2 * metric_a` with `forward()` instead of `update()`, the value returned by `.compute()` will be based only on the last batch.

This is because the default implementation of `Metric.forward()`, which is inherited by `CompositionalMetric`, calls `reset()`, which in turn resets all child metrics. Their state, however, is not restored by `forward()` afterwards.

My proposed fix would be to override `forward()` in `CompositionalMetric` and I will file a PR implementing this.

### To Reproduce

#### Code sample

```python
import torch
import torchmetrics
from torchmetrics.metric import CompositionalMetric

# Example predictions and targets
pred = torch.tensor([0.4, 0.2, 0.6, 0.7, 0.5])
target = torch.tensor([0.5, 0.4, 0.4, 0.3, 0.6])

# Compute base metric in batches
metric = torchmetrics.MeanSquaredError()
metric(preds=pred[:3], target=target[:3])
metric(preds=pred[3:], target=target[3:])
val = metric.compute()
print(val)

# Compute compositional metric in batches
cm = 1 * torchmetrics.MeanSquaredError()
cm(preds=pred[:3], target=target[:3])
cm(preds=pred[3:], target=target[3:])
val = cm.compute()
print(val)
```

### Expected behavior

Two identical values are printed:

```
tensor(0.0520)
tensor(0.0520)
```

### Actual (faulty) behavior

The compositional metric only returns the value for the last batch:

```
tensor(0.0520)
tensor(0.0850)
```

### Environment

- PyTorch Version (e.g., 1.0): 1.10.0
- OS (e.g., Linux): Linux
- How you installed PyTorch (`conda`, `pip`, source): conda
- Build command you used (if compiling from source):
- Python version: 3.8
- CUDA/cuDNN version: 11.1
- GPU models and configuration:
- Any other relevant information:

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:

Resolving patch:
diff --git a/torchmetrics/metric.py b/torchmetrics/metric.py
--- a/torchmetrics/metric.py
+++ b/torchmetrics/metric.py
@@ -756,6 +756,35 @@ def compute(self) -> Any:
 
         return self.op(val_a, val_b)
 
+    @torch.jit.unused
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+
+        val_a = (
+            self.metric_a(*args, **self.metric_a._filter_kwargs(**kwargs))
+            if isinstance(self.metric_a, Metric)
+            else self.metric_a
+        )
+        val_b = (
+            self.metric_b(*args, **self.metric_b._filter_kwargs(**kwargs))
+            if isinstance(self.metric_b, Metric)
+            else self.metric_b
+        )
+
+        if val_a is None:
+            # compute_on_step of metric_a is False
+            return None
+
+        if val_b is None:
+            if isinstance(self.metric_b, Metric):
+                # compute_on_step of metric_b is False
+                return None
+
+            # Unary op
+            return self.op(val_a)
+
+        # Binary op
+        return self.op(val_a, val_b)
+
     def reset(self) -> None:
         if isinstance(self.metric_a, Metric):
             self.metric_a.reset()
@@ -774,3 +803,6 @@ def __repr__(self) -> str:
         repr_str = self.__class__.__name__ + _op_metrics
 
         return repr_str
+
+    def _wrap_compute(self, compute: Callable) -> Callable:
+        return compute


Test patch:
diff --git a/tests/bases/test_composition.py b/tests/bases/test_composition.py
--- a/tests/bases/test_composition.py
+++ b/tests/bases/test_composition.py
@@ -24,7 +24,7 @@
 class DummyMetric(Metric):
     def __init__(self, val_to_return):
         super().__init__()
-        self._num_updates = 0
+        self.add_state("_num_updates", tensor(0), dist_reduce_fx="sum")
         self._val_to_return = val_to_return
         self._update_called = True
 
@@ -34,10 +34,6 @@ def update(self, *args, **kwargs) -> None:
     def compute(self):
         return tensor(self._val_to_return)
 
-    def reset(self):
-        self._num_updates = 0
-        return super().reset()
-
 
 @pytest.mark.parametrize(
     ["second_operand", "expected_result"],
@@ -544,7 +540,7 @@ def test_metrics_getitem(value, idx, expected_result):
 
 
 def test_compositional_metrics_update():
-
+    """test update method for compositional metrics."""
     compos = DummyMetric(5) + DummyMetric(4)
 
     assert isinstance(compos, CompositionalMetric)
@@ -557,3 +553,26 @@ def test_compositional_metrics_update():
 
     assert compos.metric_a._num_updates == 3
     assert compos.metric_b._num_updates == 3
+
+
+@pytest.mark.parametrize("compute_on_step", [True, False])
+@pytest.mark.parametrize("metric_b", [4, DummyMetric(4)])
+def test_compositional_metrics_forward(compute_on_step, metric_b):
+    """test forward method of compositional metrics."""
+    metric_a = DummyMetric(5)
+    metric_a.compute_on_step = compute_on_step
+    compos = metric_a + metric_b
+
+    assert isinstance(compos, CompositionalMetric)
+    for _ in range(3):
+        val = compos()
+        assert val == 9 if compute_on_step else val is None
+
+    assert isinstance(compos.metric_a, DummyMetric)
+    assert compos.metric_a._num_updates == 3
+
+    if isinstance(metric_b, DummyMetric):
+        assert isinstance(compos.metric_b, DummyMetric)
+        assert compos.metric_b._num_updates == 3
+
+    compos.reset()