Update for Lightning-AI#2781

Diuven · Aug 3, 2020 · c14cd67 · c14cd67
1 parent 6df2018
commit c14cd67
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pytorch_lightning/metrics/functional/classification.py b/pytorch_lightning/metrics/functional/classification.py
@@ -171,11 +171,11 @@ def stat_scores_multiple_classes(
         >>> sups
         tensor([1., 0., 1., 1.])
     """
-    num_classes = get_num_classes(pred=pred, target=target, num_classes=num_classes)
-
     if pred.ndim == target.ndim + 1:
         pred = to_categorical(pred, argmax_dim=argmax_dim)
 
+    num_classes = get_num_classes(pred=pred, target=target, num_classes=num_classes)
+
     pred = pred.view((-1, )).long()
     target = target.view((-1, )).long()