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()