Update

[ghstack-poisoned]
pytorch · Dec 17, 2024 · 1a8360b · 1a8360b
1 parent 4f41d4d
commit 1a8360b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/sota-implementations/dqn/dqn_atari.py b/sota-implementations/dqn/dqn_atari.py
@@ -86,7 +86,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         delay_value=True,
     )
     loss_module.set_keys(done="end-of-life", terminated="end-of-life")
-    loss_module.make_value_estimator(gamma=cfg.loss.gamma)
+    loss_module.make_value_estimator(gamma=cfg.loss.gamma, device=device)
     target_net_updater = HardUpdate(
         loss_module, value_network_update_interval=cfg.loss.hard_update_freq
     )

diff --git a/sota-implementations/dqn/dqn_cartpole.py b/sota-implementations/dqn/dqn_cartpole.py
@@ -69,7 +69,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         loss_function="l2",
         delay_value=True,
     )
-    loss_module.make_value_estimator(gamma=cfg.loss.gamma)
+    loss_module.make_value_estimator(gamma=cfg.loss.gamma, device=device)
     loss_module = loss_module.to(device)
     target_net_updater = HardUpdate(
         loss_module, value_network_update_interval=cfg.loss.hard_update_freq