Merge branch 'master' into feat/skorch-compatible

pyg-team · Mar 13, 2024 · 1706c96 · 1706c96
2 parents 4552e91 + 96bdf12
commit 1706c96
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/benchmark/data_frame_text_benchmark.py b/benchmark/data_frame_text_benchmark.py
@@ -457,7 +457,8 @@ def main_torch(
 
     if dataset.task_type == TaskType.BINARY_CLASSIFICATION:
         out_channels = 1
-        loss_fun = BCEWithLogitsLoss()
+        label_imbalance = sum(train_tensor_frame.y) / len(train_tensor_frame.y)
+        loss_fun = BCEWithLogitsLoss(pos_weight=1 / label_imbalance)
         metric_computer = AUROC(task='binary').to(device)
         higher_is_better = True
     elif dataset.task_type == TaskType.MULTICLASS_CLASSIFICATION: