admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-22 17:32:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

small fixes

Browse Source
This commit is contained in:
Pepijn
2025-10-14 10:46:19 +02:00
parent 2bc154e706
commit 6486982ab4
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/lerobot_train.py
+3 -1
View File
@@ -294,7 +294,7 @@ def train(cfg: TrainPipelineConfig, accelerator: Callable | None = None):
sampler=sampler, sampler=sampler,
pin_memory=device.type == "cuda", pin_memory=device.type == "cuda",
drop_last=False, drop_last=False,
prefetch_factor=2, prefetch_factor=2 if cfg.num_workers > 0 else None,
) )
if accelerator: if accelerator:
policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( 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), policy=policy if not accelerator else accelerator.unwrap_model(policy),
optimizer=optimizer, optimizer=optimizer,
scheduler=lr_scheduler, scheduler=lr_scheduler,
preprocessor=preprocessor,
postprocessor=postprocessor,
) )
update_last_checkpoint(checkpoint_dir) update_last_checkpoint(checkpoint_dir)
if wandb_logger: if wandb_logger: