Add run on cpu-only compatibility

2024-03-03 12:47:26 +01:00
parent 661bda45ea
commit b33ec5a630
5 changed files with 100 additions and 95 deletions
--- a/lerobot/common/policies/factory.py
+++ b/lerobot/common/policies/factory.py
@@ -2,7 +2,7 @@ def make_policy(cfg):
    if cfg.policy.name == "tdmpc":
        from lerobot.common.policies.tdmpc import TDMPC

-        policy = TDMPC(cfg.policy)
+        policy = TDMPC(cfg.policy, cfg.device)
    elif cfg.policy.name == "diffusion":
        from lerobot.common.policies.diffusion.policy import DiffusionPolicy