From da78460b655336d85908b8cdf417e0c5580ffc75 Mon Sep 17 00:00:00 2001
From: Pepijn <pepijn@huggingface.co>
Date: Tue, 14 Oct 2025 14:01:51 +0200
Subject: [PATCH] fix OOM bug

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

diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py
index 5b73b726c..edb88a2fb 100644
--- a/src/lerobot/scripts/lerobot_train.py
+++ b/src/lerobot/scripts/lerobot_train.py
@@ -216,7 +216,10 @@ def train(cfg: TrainPipelineConfig, accelerator: Callable | None = None):
         cfg=cfg.policy,
         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
     if accelerator: