Fix: Respect policy.device=cpu config in training

[michel-aractingi]

Type / Scope

• Type: Bug
• Scope: lerobot_train.py

Summary / Motivation

Accelerate in lerobot_train.py automatically detects and uses the best available device (CUDA or MPS > CPU), ignoring the user's policy.device setting. Trace:

1- If the user sets --policy.device=cpu
2- Preprocessor moves input batches to CPU (respecting config)
3- accelerator.prepare() moves model to CUDA (ignoring config)
4- Forward pass fails: model on CUDA, data on CPU

Reproduce

Run from main :

python -m lerobot.scripts.lerobot_train --dataset.repo_id=lerobot/pusht --policy.type=act  --policy.device=cpu  --policy.repo_id=aractingi/act_test --log_freq=1

Result: Device mismatch error between model and data.

Fix

lerobot_train.py: Force Accelerator to use cpu when the user specifies it.

force_cpu = cfg.policy.device == "cpu"
accelerator = Accelerator(..., cpu=force_cpu)

How was this tested

Run the training command with cpu and cuda or mps.

How to run locally (reviewer)

• Run the relevant tests:

  pytest -q tests/ -k <keyword>

• Run a quick example or CLI (if applicable):

  lerobot-train --some.option=true

Checklist (required before merge)

• Linting/formatting run (pre-commit run -a)
• All tests pass locally (pytest)
• Documentation updated
• CI is green

Reviewer notes

• Anything the reviewer should focus on (performance, edge-cases, specific files) or general notes.
• Anyone in the community is free to review the PR.

[Copilot]

Pull request overview

This PR fixes a bug where the Accelerator in lerobot_train.py automatically detects and uses the best available device (CUDA/MPS over CPU), ignoring the user's --policy.device=cpu configuration. This caused device mismatch errors where the model was moved to CUDA while data remained on CPU.

Changes:

• Add logic to force Accelerator to use CPU when policy.device is explicitly set to "cpu"

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[jadechoghari]

lgtm