amend

pytorch · Jul 9, 2024 · c010e39 · c010e39
1 parent 067b560
commit c010e39
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/sota-implementations/crossq/crossq.py b/sota-implementations/crossq/crossq.py
@@ -72,7 +72,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     loss_module = make_loss_module(cfg, model)
 
     # Create off-policy collector
-    collector = make_collector(cfg, train_env, exploration_policy.eval())
+    collector = make_collector(cfg, train_env, exploration_policy.eval(), device=device)
 
     # Create replay buffer
     replay_buffer = make_replay_buffer(

diff --git a/sota-implementations/crossq/utils.py b/sota-implementations/crossq/utils.py
@@ -90,15 +90,15 @@ def make_environment(cfg):
 # ---------------------------
 
 
-def make_collector(cfg, train_env, actor_model_explore):
+def make_collector(cfg, train_env, actor_model_explore, device):
     """Make collector."""
     collector = SyncDataCollector(
         train_env,
         actor_model_explore,
         init_random_frames=cfg.collector.init_random_frames,
         frames_per_batch=cfg.collector.frames_per_batch,
         total_frames=cfg.collector.total_frames,
-        device=cfg.collector.device,
+        device=device,
     )
     collector.set_seed(cfg.env.seed)
     return collector