Fix normalization of confusion matrices for empty rows

Lightning-AI · Jul 27, 2020 · 8b8b635 · 8b8b635
1 parent 3f2c102
commit 8b8b635
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pytorch_lightning/metrics/functional/classification.py b/pytorch_lightning/metrics/functional/classification.py
@@ -266,7 +266,9 @@ def confusion_matrix(
     cm = bins.reshape(num_classes, num_classes).squeeze().float()
 
     if normalize:
-        cm = cm / cm.sum(-1)
+        row_sum = cm.sum(-1, keepdim=True)
+        divisor = torch.max(row_sum, torch.tensor(1.0, device=cm.device))
+        cm = cm / divisor
 
     return cm