Fix consistency in the output of scalar tensors

CHANGELOG.md

@@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Scalar metrics will now consistently have additional dimensions squeezed ([#622](https://github.com/PyTorchLightning/metrics/pull/622))
+
 
 ### Deprecated
 

torchmetrics/aggregation.py

@@ -107,7 +107,7 @@ def update(self, value: Union[float, Tensor]) -> None:  # type: ignore
 
     def compute(self) -> Tensor:
         """Compute the aggregated value."""
-        return self.value.squeeze() if isinstance(self.value, Tensor) else self.value
+        return self.value
 
 
 class MaxMetric(BaseAggregator):
@@ -398,7 +398,7 @@ class MeanMetric(BaseAggregator):
         >>> metric.update(1)
         >>> metric.update(torch.tensor([2, 3]))
         >>> metric.compute()
-        tensor([2.])
+        tensor(2.)
     """
 
     def __init__(

torchmetrics/classification/precision_recall_curve.py

@@ -78,7 +78,7 @@ class PrecisionRecallCurve(Metric):
         >>> recall
         [tensor([1., 0.]), tensor([1., 0.]), tensor([1., 0., 0.]), tensor([1., 0., 0.]), tensor([nan, 0.])]
         >>> thresholds
-        [tensor([0.7500]), tensor([0.7500]), tensor([0.0500, 0.7500]), tensor([0.0500, 0.7500]), tensor([0.0500])]
+        [tensor(0.7500), tensor(0.7500), tensor([0.0500, 0.7500]), tensor([0.0500, 0.7500]), tensor(0.0500)]
     """
 
     is_differentiable = False

torchmetrics/image/lpip_similarity.py

@@ -84,7 +84,7 @@ class LPIPS(Metric):
         >>> img1 = torch.rand(10, 3, 100, 100)
         >>> img2 = torch.rand(10, 3, 100, 100)
         >>> lpips(img1, img2)
-        tensor([0.3566], grad_fn=<DivBackward0>)
+        tensor(0.3566, grad_fn=<SqueezeBackward0>)
     """
 
     is_differentiable = True

torchmetrics/metric.py

@@ -25,7 +25,15 @@
 from torch.nn import Module
 
 from torchmetrics.utilities import apply_to_collection, rank_zero_warn
-from torchmetrics.utilities.data import _flatten, dim_zero_cat, dim_zero_max, dim_zero_mean, dim_zero_min, dim_zero_sum
+from torchmetrics.utilities.data import (
+    _flatten,
+    _squeeze_if_scalar,
+    dim_zero_cat,
+    dim_zero_max,
+    dim_zero_mean,
+    dim_zero_min,
+    dim_zero_sum,
+)
 from torchmetrics.utilities.distributed import gather_all_tensors
 from torchmetrics.utilities.exceptions import TorchMetricsUserError
 from torchmetrics.utilities.imports import _LIGHTNING_AVAILABLE, _compare_version
@@ -369,7 +377,8 @@ def wrapped_func(*args: Any, **kwargs: Any) -> Any:
             with self.sync_context(
                 dist_sync_fn=self.dist_sync_fn, should_sync=self._to_sync, should_unsync=self._should_unsync
             ):
-                self._computed = compute(*args, **kwargs)
+                value = compute(*args, **kwargs)
+                self._computed = _squeeze_if_scalar(value)
 
             return self._computed
 

torchmetrics/utilities/data.py

@@ -238,3 +238,11 @@ def get_group_indexes(indexes: Tensor) -> List[Tensor]:
             res[_id] = [i]
 
     return [tensor(x, dtype=torch.long) for x in res.values()]
+
+
+def _squeeze_scalar_element_tensor(x: Tensor) -> Tensor:
+    return x.squeeze() if x.numel() == 1 else x
+
+
+def _squeeze_if_scalar(data: Any) -> Any:
+    return apply_to_collection(data, Tensor, _squeeze_scalar_element_tensor)