import

Lightning-Universe · Jul 2, 2020 · 1209662 · 1209662
1 parent 9552d3a
commit 1209662
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/pl_bolts/loggers/__init__.py b/pl_bolts/loggers/__init__.py
@@ -5,7 +5,7 @@
 __all__ = []
 
 try:
-    from pytorch_lightning.loggers.trains import TrainsLogger
+    from pl_bolts.loggers.trains import TrainsLogger
 except ImportError:  # pragma: no-cover
     pass  # pragma: no-cover
 else:

diff --git a/tests/loggers/__init__.py b/tests/loggers/__init__.py
@@ -0,0 +1,8 @@
+__all__ = []
+
+try:
+    from pl_bolts.loggers.trains import TrainsLogger
+except ImportError:  # pragma: no-cover
+    pass  # pragma: no-cover
+else:
+    __all__.append('TrainsLogger')
diff --git a/tests/loggers/test_trains.py b/tests/loggers/test_trains.py
@@ -1,9 +1,9 @@
 import pickle
 
 from pytorch_lightning import Trainer
-from pytorch_lightning.loggers import TrainsLogger
 
 from pl_bolts import LitMNIST
+from pl_bolts.loggers import TrainsLogger
 
 
 def test_trains_logger(tmpdir):