Add "none" option to AveragePrecision (#1116)

Lightning-AI · Jun 28, 2022 · bdb3e59 · bdb3e59
1 parent d9c345b
commit bdb3e59
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -49,6 +49,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Skip box conversion if no boxes are present in `MeanAveragePrecision` ([#1097](https://github.com/Lightning-AI/metrics/pull/1097))
 
 
+- Fixed inconsistency in docs and code when setting `average="none"` in `AvaragePrecision` metric ([#1116](https://github.com/Lightning-AI/metrics/pull/1116))
+
+
 ## [0.9.1] - 2022-06-08
 
 ### Added

diff --git a/src/torchmetrics/functional/classification/average_precision.py b/src/torchmetrics/functional/classification/average_precision.py
@@ -167,9 +167,9 @@ def _average_precision_compute_with_precision_recall(
             return res[~torch.isnan(res)].mean()
         weights = torch.ones_like(res) if weights is None else weights
         return (res * weights)[~torch.isnan(res)].sum()
-    if average is None:
+    if average is None or average == "none":
         return res
-    allowed_average = ("micro", "macro", "weighted", None)
+    allowed_average = ("micro", "macro", "weighted", "none", None)
     raise ValueError(f"Expected argument `average` to be one of {allowed_average}" f" but got {average}")