From 6486982ab41bb4e285afd9a7ba5ae0b4382fb87d Mon Sep 17 00:00:00 2001
From: Pepijn <pepijn@huggingface.co>
Date: Tue, 14 Oct 2025 10:46:19 +0200
Subject: [PATCH] small fixes

---
 src/lerobot/scripts/lerobot_train.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py
index 77da7899f..88fa26997 100644
--- a/src/lerobot/scripts/lerobot_train.py
+++ b/src/lerobot/scripts/lerobot_train.py
@@ -294,7 +294,7 @@ def train(cfg: TrainPipelineConfig, accelerator: Callable | None = None):
         sampler=sampler,
         pin_memory=device.type == "cuda",
         drop_last=False,
-        prefetch_factor=2,
+        prefetch_factor=2 if cfg.num_workers > 0 else None,
     )
     if accelerator:
         policy, optimizer, dataloader, lr_scheduler = accelerator.prepare(
@@ -369,6 +369,8 @@ def train(cfg: TrainPipelineConfig, accelerator: Callable | None = None):
                 policy=policy if not accelerator else accelerator.unwrap_model(policy),
                 optimizer=optimizer,
                 scheduler=lr_scheduler,
+                preprocessor=preprocessor,
+                postprocessor=postprocessor,
             )
             update_last_checkpoint(checkpoint_dir)
             if wandb_logger: