Fix incorrect caching of MetricCollection

src/torchmetrics/collections.py

@@ -206,6 +206,11 @@ def update(self, *args: Any, **kwargs: Any) -> None:
  """
  # Use compute groups if already initialized and checked
  if self._groups_checked:
+ # Delete the cache of all metrics to invalidate the cache and therefore recent compute calls, forcing new
+ # compute calls to recompute
+ for k in self.keys(keep_base=True):
+ mi = getattr(self, k)
+ mi._computed = None
  for cg in self._groups.values():
  # only update the first member
  m0 = getattr(self, cg[0])
@@ -304,7 +309,6 @@ def _compute_groups_create_state_ref(self, copy: bool = False) -> None:
  # Determine if we just should set a reference or a full copy
  setattr(mi, state, deepcopy(m0_state) if copy else m0_state)
  mi._update_count = deepcopy(m0._update_count) if copy else m0._update_count
- mi._computed = deepcopy(m0._computed) if copy else m0._computed
  self._state_is_copy = copy
 
  def compute(self) -> Dict[str, Any]:

src/torchmetrics/metric.py

@@ -619,7 +619,7 @@ def wrapped_func(*args: Any, **kwargs: Any) -> Any:
 
  # return cached value
  if self._computed is not None:
- return self._computed
+ return deepcopy(self._computed)
 
  # compute relies on the sync context manager to gather the states across processes and apply reduction
  # if synchronization happened, the current rank accumulated states will be restored to keep
@@ -634,7 +634,8 @@ def wrapped_func(*args: Any, **kwargs: Any) -> Any:
  if self.compute_with_cache:
  self._computed = value
 
- return value
+ # Return a deep copy to avoid side effects for non-scalar values, e.g. ConfusionMatrix
+ return deepcopy(value)
 
  return wrapped_func
 

tests/unittests/bases/test_collections.py

@@ -452,10 +452,29 @@ def test_check_compute_groups_correctness(self, metrics, expected, preds, target
  for key in res_cg:
  assert torch.allclose(res_cg[key], res_without_cg[key])
 
+ # Check if second compute is the same
+ res_cg2 = m.compute()
+ for key in res_cg2:
+ assert torch.allclose(res_cg[key], res_cg2[key])
+
  if with_reset:
  m.reset()
  m2.reset()
 
+ # Test if a second compute without a reset is the same
+ m.reset()
+ m.update(preds, target)
+ res_cg = m.compute()
+ # Simulate different preds by simply inversing them
+ m.update(1 - preds, target)
+ res_cg2 = m.compute()
+ # Now check if the results from the first compute are different from the second
+ for key in res_cg:
+ # A different shape is okay, therefore skip (this happens for multidim_average="samplewise")
+ if res_cg[key].shape != res_cg2[key].shape:
+ continue
+ assert not torch.all(res_cg[key] == res_cg2[key])
+
  @pytest.mark.parametrize("method", ["items", "values", "keys"])
  def test_check_compute_groups_items_and_values(self, metrics, expected, preds, target, method):
  """Check states are copied instead of passed by ref when a single metric in the collection is access."""