fix: optional tensorboard dependency

Author: samedii

lantern/__init__.py

@@ -14,7 +14,11 @@
 from lantern.set_seeds import set_seeds
 from lantern.worker_init_fn import worker_init_fn
 from lantern.progress_bar import ProgressBar
-from lantern.early_stopping import EarlyStopping
+
+try:
+    from lantern.early_stopping import EarlyStopping
+except ImportError:
+    pass
 from lantern.git_info import git_info
 
 from pkg_resources import get_distribution, DistributionNotFound

lantern/early_stopping.py

@@ -1,13 +1,20 @@
-import torch.utils.tensorboard
 from typing import Optional
 
 from lantern import FunctionalBase
 
+try:
+    import torch.utils.tensorboard
+
+    SummaryWriter = torch.utils.tensorboard.SummaryWriter
+except ImportError:
+    SummaryWriter = None
+    pass
+
 
 class EarlyStopping(FunctionalBase):
     """Keeps track of the best score and how long ago it was calculated."""
 
-    tensorboard_logger: torch.utils.tensorboard.SummaryWriter
+    tensorboard_logger: SummaryWriter
     best_score: Optional[float] = None
     scores_since_improvement: int = -1