[Feature] Add all implemented loggers to the init of loggers (pytorch…

…#402)
ai4co · Sep 9, 2022 · b8c0cde · b8c0cde
1 parent 3bdae23
commit b8c0cde
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/test/test_trainer.py b/test/test_trainer.py
@@ -13,7 +13,7 @@
 
 try:
     from tensorboard.backend.event_processing import event_accumulator
-    from torchrl.trainers.loggers import TensorboardLogger
+    from torchrl.trainers.loggers.tensorboard import TensorboardLogger
 
     _has_tb = True
 except ImportError:
@@ -242,6 +242,7 @@ def test_recorder():
         args.record_frames = 24 // args.frame_skip
         args.record_interval = 2
         args.catframes = 4
+        args.collector_devices = ["cpu"]
 
         N = 8
 

diff --git a/torchrl/trainers/loggers/__init__.py b/torchrl/trainers/loggers/__init__.py
@@ -4,3 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 from .common import Logger
+from .csv import CSVLogger
+from .tensorboard import TensorboardLogger
+from .wandb import WandbLogger