diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index e5c5d6a4c..00676a011 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -97,7 +97,7 @@ def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]: policy_key = env_cfg.features_map[key] policy_features[policy_key] = feature - + return policy_features diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 2d51a3881..ef56bdb61 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -169,7 +169,7 @@ def make_policy( else: # Make a fresh policy. policy = policy_cls(**kwargs) - + policy.to(cfg.device) assert isinstance(policy, nn.Module)