From f7283193ea9ae932423e3a1e27524a27fa5c0fe5 Mon Sep 17 00:00:00 2001 From: Adil Zouitine Date: Mon, 22 Sep 2025 11:26:30 +0200 Subject: [PATCH] fix(trainer): overrides device to the target device, for the device processor on the preprocessor (#1993) * fix(trainer): overiddes device to the target defice, for device processor on preprocessor * Update src/lerobot/scripts/train.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Adil Zouitine --------- Signed-off-by: Adil Zouitine Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/lerobot/scripts/train.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/lerobot/scripts/train.py b/src/lerobot/scripts/train.py index 485fc927..5594d2f9 100644 --- a/src/lerobot/scripts/train.py +++ b/src/lerobot/scripts/train.py @@ -183,6 +183,9 @@ def train(cfg: TrainPipelineConfig): # Only provide dataset_stats when not resuming from saved processor state processor_kwargs["dataset_stats"] = dataset.meta.stats + if cfg.policy.pretrained_path is not None: + processor_kwargs["preprocessor_overrides"] = {"device_processor": {"device": device.type}} + preprocessor, postprocessor = make_pre_post_processors( policy_cfg=cfg.policy, pretrained_path=cfg.policy.pretrained_path, **processor_kwargs )