Update

[ghstack-poisoned]

Author: vmoens

torchrl/collectors/_single.py

@@ -42,8 +42,8 @@
     set_exploration_type,
 )
 from torchrl.modules import RandomPolicy, set_exploration_modules_spec_from_env
-from torchrl.weight_update import WeightSyncScheme
 from torchrl.weight_update.utils import _resolve_model
+from torchrl.weight_update.weight_sync_schemes import WeightSyncScheme
 
 
 class _CollectorProfiler:

torchrl/weight_update/_shared.py

@@ -10,8 +10,7 @@
 
 from torch import multiprocessing as mp, nn
 
-from torchrl._utils import logger as torchrl_logger
-from torchrl.collectors._constants import WEIGHT_SYNC_TIMEOUT
+from torchrl._utils import logger as torchrl_logger, WEIGHT_SYNC_TIMEOUT
 
 from torchrl.weight_update.utils import _resolve_model
 from torchrl.weight_update.weight_sync_schemes import (