fix ignored custom callable in retrieval metric aggregation (Lightning-AI#2364)

* fix retrieval aggregation

* fix retrieval tests

* changelog

---------

Co-authored-by: Nicki Skafte Detlefsen <skaftenicki@gmail.com>

Author: fschlatt

CHANGELOG.md

@@ -36,6 +36,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed cached network in `FeatureShare` not being moved to the correct device ([#2348](https://github.com/Lightning-AI/torchmetrics/pull/2348))
 
 
+- Fixed custom aggregation in retrieval metrics ([#2364](https://github.com/Lightning-AI/torchmetrics/pull/2364))
+
+
 - Fixed initialize aggregation metrics with default floating type ([#2366](https://github.com/Lightning-AI/torchmetrics/pull/2366))
 
 ---

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)
 

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(