fix: adjust BCEWithLogitsLoss to handle class imbalance with pos_weight

app/training/train_deep_classifiers.py

@@ -209,7 +209,11 @@ def train_one_fold(
     scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
         optimizer, mode="max", factor=0.5, patience=5
     )
-    criterion = nn.BCEWithLogitsLoss()
+    # pos_weight compensates for class imbalance (n_neg / n_pos)
+    n_pos = max(int(y_train.sum()), 1)
+    n_neg = len(y_train) - n_pos
+    pos_weight = torch.tensor([n_neg / n_pos], dtype=torch.float32).to(DEVICE)
+    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
 
     best_auc = 0.0
     best_probs = None