Add argument comments

Lightning-AI · Jul 27, 2020 · 913c1e3 · 913c1e3
1 parent 83927a9
commit 913c1e3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/pytorch_lightning/metrics/classification.py b/pytorch_lightning/metrics/classification.py
@@ -103,6 +103,7 @@ def __init__(
             normalize: whether to compute a normalized confusion matrix
             reduce_group: the process group to reduce metric results from DDP
             reduce_op: the operation to perform for ddp reduction
+            num_classes: number of classes if known. Important for DDP reduction.
         """
         super().__init__(name='confusion_matrix',
                          reduce_group=reduce_group,

diff --git a/pytorch_lightning/metrics/functional/classification.py b/pytorch_lightning/metrics/functional/classification.py
@@ -245,6 +245,7 @@ def confusion_matrix(
         pred: estimated targets
         target: ground truth labels
         normalize: normalizes confusion matrix
+        num_classes: number of classes if known
 
     Return:
         Tensor, confusion matrix C [num_classes, num_classes ]