add wandb logger to main module

webis-de · Nov 21, 2024 · e84281c · e84281c
1 parent 659c172
commit e84281c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lightning_ir/__init__.py b/lightning_ir/__init__.py
@@ -68,7 +68,7 @@
     RankNet,
     SupervisedMarginMSE,
 )
-from .main import LightningIRTrainer
+from .main import LightningIRTrainer, LightningIRWandbLogger
 from .models import (
     ColConfig,
     ColModel,
@@ -177,6 +177,7 @@
     "LightningIRTokenizer",
     "LightningIRTokenizerClassFactory",
     "LightningIRTrainer",
+    "LightningIRWandbLogger",
     "LinearLRSchedulerWithLinearWarmup",
     "LocalizedContrastiveEstimation",
     "LR_SCHEDULERS",