Apply mask for labels for multi-class recall

Summary: Fixes pytorch#150 Differential Revision: D47241862 fbshipit-source-id: 2e76521caefceb1b8ac8a355d4748347d6432403
ananthsub · Jul 5, 2023 · 873446f · 873446f
1 parent c7a6ba9
commit 873446f
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torcheval/metrics/functional/classification/recall.py b/torcheval/metrics/functional/classification/recall.py
@@ -191,6 +191,7 @@ def _recall_compute(
         # Ignore classes which have no samples in `target` and `input`
         mask = (num_labels != 0) | (num_predictions != 0)
         num_tp = num_tp[mask]
+        num_labels = num_labels[mask]
 
     recall = num_tp / num_labels