Update

[ghstack-poisoned]
pytorch · Dec 17, 2024 · 7444057 · 7444057
1 parent 7c044b3
commit 7444057
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/sota-implementations/decision_transformer/dt.py b/sota-implementations/decision_transformer/dt.py
@@ -76,7 +76,9 @@ def main(cfg: "DictConfig"):  # noqa: F821
     loss_module = make_dt_loss(cfg.loss, actor, device=model_device)
 
     # Create optimizer
-    transformer_optim, scheduler = make_dt_optimizer(cfg.optim, loss_module)
+    transformer_optim, scheduler = make_dt_optimizer(
+        cfg.optim, loss_module, model_device
+    )
 
     # Create inference policy
     inference_policy = DecisionTransformerInferenceWrapper(

diff --git a/sota-implementations/decision_transformer/utils.py b/sota-implementations/decision_transformer/utils.py
@@ -511,10 +511,10 @@ def make_odt_optimizer(optim_cfg, loss_module):
     return dt_optimizer, log_temp_optimizer, scheduler
 
 
-def make_dt_optimizer(optim_cfg, loss_module):
+def make_dt_optimizer(optim_cfg, loss_module, device):
     dt_optimizer = torch.optim.Adam(
         loss_module.actor_network_params.flatten_keys().values(),
-        lr=torch.as_tensor(optim_cfg.lr),
+        lr=torch.tensor(optim_cfg.lr, device=device),
         weight_decay=optim_cfg.weight_decay,
         eps=1.0e-8,
     )

diff --git a/sota-implementations/dqn/config_atari.yaml b/sota-implementations/dqn/config_atari.yaml
@@ -42,5 +42,5 @@ loss:
 
 compile:
   compile: False
-  compile_mode:
+  compile_mode: default
   cudagraphs: False