Update

[ghstack-poisoned]
pytorch · Dec 14, 2024 · 105440a · 105440a
2 parents 5fa2d11 + 1013883
commit 105440a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sota-implementations/decision_transformer/utils.py b/sota-implementations/decision_transformer/utils.py
@@ -395,7 +395,7 @@ def make_odt_model(cfg, device: torch.device | None = None) -> TensorDictModule:
     with torch.no_grad(), set_exploration_type(ExplorationType.RANDOM):
         td = proof_environment.rollout(max_steps=100)
         td["action"] = td["next", "action"]
-        actor(td)
+        actor(td.to(device))
 
     return actor