diff options
Diffstat (limited to 'training')
-rw-r--r-- | training/run_experiment.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/training/run_experiment.py b/training/run_experiment.py index 1fb4bc2..491112c 100644 --- a/training/run_experiment.py +++ b/training/run_experiment.py @@ -69,7 +69,7 @@ def _configure_callbacks( def _configure_logger( network: Type[nn.Module], args: Dict, use_wandb: bool -) -> pl.loggers.WandbLogger: +) -> Type[pl.loggers.LightningLoggerBase]: """Configures lightning logger.""" if use_wandb: pl_logger = pl.loggers.WandbLogger() |