Add None as reduction option in CosineSimilarity

CHANGELOG.md

@@ -20,6 +20,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added Symmetric Mean Absolute Percentage error (SMAPE) ([#375](https://github.com/PyTorchLightning/metrics/issues/375))
 
 
+- Added `None` as reduction option in `CosineSimilarity` metric ([#400](https://github.com/PyTorchLightning/metrics/pull/400))
+
+
 - Added Permutation Invariant Training metric (PIT) ([#294](https://github.com/PyTorchLightning/metrics/issues/294))
 
 

torchmetrics/functional/regression/cosine_similarity.py

@@ -61,7 +61,12 @@ def _cosine_similarity_compute(preds: Tensor, target: Tensor, reduction: str = '
     preds_norm = preds.norm(dim=-1)
     target_norm = target.norm(dim=-1)
     similarity = dot_product / (preds_norm * target_norm)
-    reduction_mapping = {"sum": torch.sum, "mean": torch.mean, "none": lambda x: x}
+    reduction_mapping = {
+        "sum": torch.sum,
+        "mean": torch.mean,
+        "none": lambda x: x,
+        None: lambda x: x,
+    }
     return reduction_mapping[reduction](similarity)
 
 

torchmetrics/regression/cosine_similarity.py

@@ -77,10 +77,13 @@ def __init__(
             process_group=process_group,
             dist_sync_fn=dist_sync_fn
         )
+        allowed_reduction = ('sum', 'mean', 'none', None)
+        if reduction not in allowed_reduction:
+            raise ValueError(f"Expected argument `reduction` to be one of {allowed_reduction}" f" but got {reduction}")
+        self.reduction = reduction
 
         self.add_state("preds", [], dist_reduce_fx="cat")
         self.add_state("target", [], dist_reduce_fx="cat")
-        self.reduction = reduction
 
     def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
         """