[Feature] Trainer Infrastructure - Timing and Utilities

ghstack-source-id: aa87e8c
Pull-Request: #3188

Author: vmoens

test/_utils_internal.py

@@ -21,13 +21,7 @@
 from tensordict.nn import TensorDictModuleBase
 from torch import nn, vmap
 
-from torchrl._utils import (
-    implement_for,
-    logger,
-    logger as torchrl_logger,
-    RL_WARNINGS,
-    seed_generator,
-)
+from torchrl._utils import implement_for, logger, RL_WARNINGS, seed_generator
 from torchrl.data.utils import CloudpickleWrapper
 from torchrl.envs import MultiThreadedEnv, ObservationNorm
 from torchrl.envs.batched_envs import ParallelEnv, SerialEnv
@@ -230,7 +224,7 @@ def f_retry(*args, **kwargs):
                     return f(*args, **kwargs)
                 except ExceptionToCheck as e:
                     msg = "%s, Retrying in %d seconds..." % (str(e), mdelay)
-                    torchrl_logger.info(msg)
+                    logger.info(msg)
                     time.sleep(mdelay)
                     mtries -= 1
             try:

torchrl/__init__.py

@@ -62,7 +62,7 @@
     timeit,
 )
 
-torchrl_logger = logger
+logger = logger
 
 # Filter warnings in subprocesses: True by default given the multiple optional
 # deps of the library. This can be turned on via `torchrl.filter_warnings_subprocess = False`.
@@ -121,5 +121,5 @@ def _inv(self):
     "set_auto_unwrap_transformed_env",
     "timeit",
     "logger",
-    "torchrl_logger",
+    "logger",
 ]