From 91ff9c4975351d39ec3b9659da654bf8bd895cfe Mon Sep 17 00:00:00 2001 From: Michel Aractingi Date: Mon, 12 Jan 2026 12:19:02 +0100 Subject: [PATCH] Fix: Respect policy.device=cpu config in training (#2778) * fix cpu training in lerobot_train * Update src/lerobot/scripts/lerobot_train.py Signed-off-by: Michel Aractingi --- src/lerobot/scripts/lerobot_train.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 286c69906..927e9659a 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -259,7 +259,14 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None): from accelerate.utils import DistributedDataParallelKwargs ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) - accelerator = Accelerator(step_scheduler_with_optimizer=False, kwargs_handlers=[ddp_kwargs]) + # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting. + # Force the device to be CPU when policy.device is set to CPU. + force_cpu = cfg.policy.device == "cpu" + accelerator = Accelerator( + step_scheduler_with_optimizer=False, + kwargs_handlers=[ddp_kwargs], + cpu=force_cpu, + ) init_logging(accelerator=accelerator)