Repository: Lightning-AI/torchmetrics

Issue:
Custom callable is ignored in retrieval metric aggregation
## 🐛 Bug

When passing a custom aggregation to any `RetrievalMetric`, the callable is ignored. There is a bug in the retrieval aggregation function that defaults to using the maximum of the values by default https://github.com/Lightning-AI/torchmetrics/blob/1a2f406da17a01c13ce1add83835e097460cbce3/src/torchmetrics/retrieval/base.py#L38

### To Reproduce

The following custom aggregation should output the tensor of values without aggregation but returns 1, the max of the values.

```
>>> from torchmetrics.retrieval import RetrievalNormalizedDCG
>>> indexes = tensor([0, 0, 0, 1, 1, 1, 1])
>>> preds = tensor([0.2, 0.3, 0.5, 0.1, 0.3, 0.5, 0.2])
>>> target = tensor([False, False, True, False, True, False, True])
>>> ndcg = RetrievalNormalizedDCG(aggregation=lambda x, dim: x)
>>> print(ndcg(preds, target, indexes=indexes))
tensor(1.)
>>> ndcg = RetrievalNormalizedDCG(aggregation="max")
>>> print(ndcg(preds, target, indexes=indexes))
tensor(1.)
>>> ndcg = RetrievalNormalizedDCG()
>>> print(ndcg(preds, target, indexes=indexes))
tensor(0.8467)
```


Discussion:
Hi! thanks for your contribution!, great first issue!

Resolving patch:
diff --git a/src/torchmetrics/retrieval/base.py b/src/torchmetrics/retrieval/base.py
--- a/src/torchmetrics/retrieval/base.py
+++ b/src/torchmetrics/retrieval/base.py
@@ -35,7 +35,7 @@ def _retrieval_aggregate(
         return values.median() if dim is None else values.median(dim=dim).values
     if aggregation == "min":
         return values.min() if dim is None else values.min(dim=dim).values
-    if aggregation:
+    if aggregation == "max":
         return values.max() if dim is None else values.max(dim=dim).values
     return aggregation(values, dim=dim)
 


Test patch:
diff --git a/tests/unittests/retrieval/helpers.py b/tests/unittests/retrieval/helpers.py
--- a/tests/unittests/retrieval/helpers.py
+++ b/tests/unittests/retrieval/helpers.py
@@ -19,7 +19,6 @@
 import torch
 from numpy import array
 from torch import Tensor, tensor
-from torchmetrics.retrieval.base import _retrieval_aggregate
 from typing_extensions import Literal
 
 from unittests.helpers import seed_all
@@ -42,6 +41,23 @@
 # a version of get_group_indexes that depends on NumPy is here to avoid this dependency for the full library
 
 
+def _retrieval_aggregate(
+    values: Tensor,
+    aggregation: Union[Literal["mean", "median", "min", "max"], Callable] = "mean",
+    dim: Optional[int] = None,
+) -> Tensor:
+    """Aggregate the final retrieval values into a single value."""
+    if aggregation == "mean":
+        return values.mean() if dim is None else values.mean(dim=dim)
+    if aggregation == "median":
+        return values.median() if dim is None else values.median(dim=dim).values
+    if aggregation == "min":
+        return values.min() if dim is None else values.min(dim=dim).values
+    if aggregation == "max":
+        return values.max() if dim is None else values.max(dim=dim).values
+    return aggregation(values, dim=dim)
+
+
 def get_group_indexes(indexes: Union[Tensor, np.ndarray]) -> List[Union[Tensor, np.ndarray]]:
     """Extract group indexes.
 
@@ -74,7 +90,7 @@ def get_group_indexes(indexes: Union[Tensor, np.ndarray]) -> List[Union[Tensor,
 
 
 def _custom_aggregate_fn(val: Tensor, dim=None) -> Tensor:
-    return (val**2).mean(dim=dim)
+    return (val**2).mean() if dim is None else (val**2).mean(dim=dim)
 
 
 def _compute_sklearn_metric(