fix row-wise division while normalize

Lightning-AI · Sep 12, 2020 · fd19c78 · fd19c78
1 parent 3ba5e19
commit fd19c78
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/metrics/functional/classification.py b/pytorch_lightning/metrics/functional/classification.py
@@ -312,7 +312,7 @@ def confusion_matrix(
     cm = bins.reshape(num_classes, num_classes).squeeze().float()
 
     if normalize:
-        cm = cm / cm.sum(-1)
+        cm = cm / cm.sum(-1, keepdim=True)
         nan_elements = cm[torch.isnan(cm)].nelement()
         if nan_elements != 0:
             cm[torch.isnan(cm)] = 0