Repository: Lightning-AI/torchmetrics

Issue:
Usage of MetricCollection throws warnings that compute method was called before update.
## 🐛 Bug

MetricCollection is using groups now. For e.g. Precision(), Recall(), F1Score() are put in a same group. These groups aren't being updated. Warnings are issued, but I've not checked whether the results are calculated correctly. 

Warning: 
```
torchmetrics/utilities/prints.py:36: UserWarning: The ``compute`` method of metric Precision was called before the ``update`` method which may lead to errors, as metric states have not yet been updated.
torchmetrics/utilities/prints.py:36: UserWarning: The ``compute`` method of metric Recall was called before the ``update`` method which may lead to errors, as metric states have not yet been updated.
torchmetrics/utilities/prints.py:36: UserWarning: The ``compute`` method of metric F1Score was called before the ``update`` method which may lead to errors, as metric states have not yet been updated.
```
### To Reproduce

Run following code:

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

#### Code sample
1. Initialize MetricCollection
```python
# initalize
import torchmetrics
import torch

metrics = torchmetrics.MetricCollection([
    torchmetrics.Accuracy(),
    torchmetrics.Precision(num_classes=3),
    torchmetrics.Recall(num_classes=3),
    torchmetrics.F1Score(num_classes=3),
    torchmetrics.ConfusionMatrix(num_classes=3)
])
```

2. Simulate first epoch:
```python
# epoch 1
target = torch.tensor([0, 1, 2, 0, 1, 2])
preds = torch.tensor([0, 2, 1, 0, 0, 1])
metrics.update(preds, target)
metrics.update(preds, target)
metrics.update(preds, target)

for k, m in metrics.items():
    print("Update called flag of ", k, ":", m._update_called)

print(metrics.compute())
metrics.reset()
```

3. Simulate second epoch, where the warnings occur:
```python
# epoch 2
preds = torch.tensor([0, 1, 2, 0, 1, 2])
target = torch.tensor([0, 2, 1, 0, 0, 1])
metrics.update(preds, target)
metrics.update(preds, target)
metrics.update(preds, target)

for k, m in metrics.items():
    print("Update called flag of ", k, ":", m._update_called) # <--- will return false for precision, recall and f1score 

print(metrics.compute())  # <--- warning occurs here
print(metrics.reset())
``` 

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): 0.8.2, pip
- Python & PyTorch Version (e.g., 1.0): python 3.8.6 and torch 1.9.1 
- Any other relevant information such as OS (e.g., Linux): Linux



Discussion:

Resolving patch:
diff --git a/torchmetrics/collections.py b/torchmetrics/collections.py
--- a/torchmetrics/collections.py
+++ b/torchmetrics/collections.py
@@ -169,6 +169,9 @@ def update(self, *args: Any, **kwargs: Any) -> None:
                 # only update the first member
                 m0 = getattr(self, cg[0])
                 m0.update(*args, **m0._filter_kwargs(**kwargs))
+                for i in range(1, len(cg)):  # copy over the update count
+                    mi = getattr(self, cg[i])
+                    mi._update_count = m0._update_count
         else:  # the first update always do per metric to form compute groups
             for _, m in self.items(keep_base=True):
                 m_kwargs = m._filter_kwargs(**kwargs)


Test patch:
diff --git a/tests/bases/test_collections.py b/tests/bases/test_collections.py
--- a/tests/bases/test_collections.py
+++ b/tests/bases/test_collections.py
@@ -340,25 +340,35 @@ def test_check_compute_groups(metrics, expected, prefix, postfix):
     assert len(m.compute_groups) == len(m)
     assert m2.compute_groups == {}
 
-    preds = torch.randn(10, 3).softmax(dim=-1)
-    target = torch.randint(3, (10,))
-    m.update(preds, target)
-    m2.update(preds, target)
+    for _ in range(2):  # repeat to emulate effect of multiple epochs
+        preds = torch.randn(10, 3).softmax(dim=-1)
+        target = torch.randint(3, (10,))
+        m.update(preds, target)
+        m2.update(preds, target)
 
-    assert m.compute_groups == expected
-    assert m2.compute_groups == {}
+        for _, member in m.items():
+            assert member._update_called
 
-    preds = torch.randn(10, 3).softmax(dim=-1)
-    target = torch.randint(3, (10,))
-    # compute groups should kick in here
-    m.update(preds, target)
-    m2.update(preds, target)
+        assert m.compute_groups == expected
+        assert m2.compute_groups == {}
+
+        preds = torch.randn(10, 3).softmax(dim=-1)
+        target = torch.randint(3, (10,))
+        # compute groups should kick in here
+        m.update(preds, target)
+        m2.update(preds, target)
+
+        for _, member in m.items():
+            assert member._update_called
+
+        # compare results for correctness
+        res_cg = m.compute()
+        res_without_cg = m2.compute()
+        for key in res_cg.keys():
+            assert torch.allclose(res_cg[key], res_without_cg[key])
 
-    # compare results for correctness
-    res_cg = m.compute()
-    res_without_cg = m2.compute()
-    for key in res_cg.keys():
-        assert torch.allclose(res_cg[key], res_without_cg[key])
+        m.reset()
+        m2.reset()
 
 
 @pytest.mark.parametrize(