Fix division by zero error for labels with no samples

fcakyon · Nov 22, 2024 · 98b2612 · 98b2612
1 parent 1c71f91
commit 98b2612
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/balanced_loss/losses.py b/balanced_loss/losses.py
@@ -97,6 +97,8 @@ def forward(
 
         if self.class_balanced:
             effective_num = 1.0 - np.power(self.beta, self.samples_per_class)
+            # Avoid division by 0 error for test cases without all labels present.
+            effective_num[effective_num == 0] = 1
             weights = (1.0 - self.beta) / np.array(effective_num)
             weights = weights / np.sum(weights) * num_classes
             weights = torch.tensor(weights, device=logits.device).float()