From 6f662ee17e47f132a32aa43e7a113fca9851e0f3 Mon Sep 17 00:00:00 2001
From: Marcus Campbell <21266489+marcus-campbell@users.noreply.github.com>
Date: Fri, 21 Aug 2026 15:41:16 -0700
Subject: [PATCH] Fix BiC bias-correction loss

Pass the bias-corrected logits directly to F.cross_entropy during
stage-two training. F.cross_entropy already applies log-softmax
internally, so normalizing the logits first changes the loss and
gradients. This fix matches the loss defined by the original BiC paper
and reference implementation.
---
 models/bic.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bic.py b/models/bic.py
index c57aba6..af5531e 100644
--- a/models/bic.py
+++ b/models/bic.py
@@ -115,7 +115,7 @@ def _run(self, train_loader, test_loader, optimizer, scheduler, stage):
                     else:
                         loss = clf_loss
                 elif stage == "bias_correction":
-                    loss = F.cross_entropy(torch.softmax(logits, dim=1), targets)
+                    loss = F.cross_entropy(logits, targets)
                 else:
                     raise NotImplementedError()