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

fix OOM bug

Browse Source
This commit is contained in:
Pepijn
2025-10-14 14:01:51 +02:00
parent a0d0b00e04
commit da78460b65
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/lerobot_train.py
+4 -1
View File
@@ -216,7 +216,10 @@ def train(cfg: TrainPipelineConfig, accelerator: Callable | None = None):
cfg=cfg.policy, cfg=cfg.policy,
ds_meta=dataset.meta, ds_meta=dataset.meta,
) )
policy.to(device)
# Only move to device if not using accelerator (accelerator.prepare will handle device placement)
if not accelerator:
policy.to(device)
# Wait for all processes to finish policy creation before continuing # Wait for all processes to finish policy creation before continuing
if accelerator: if accelerator: