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

fix accel

Browse Source
This commit is contained in:
Pepijn
2025-09-23 22:32:22 +02:00
parent b794fc3c70
commit fc7998a3d5
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/train.py
+4
View File
@@ -21,6 +21,7 @@ from typing import Any
import torch import torch
from accelerate import Accelerator from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs
from termcolor import colored from termcolor import colored
from torch.amp import GradScaler from torch.amp import GradScaler
from torch.optim import Optimizer from torch.optim import Optimizer
@@ -176,9 +177,12 @@ def train(cfg: TrainPipelineConfig):
# Initialize Accelerate if requested # Initialize Accelerate if requested
accelerator = None accelerator = None
if cfg.use_accelerate: if cfg.use_accelerate:
# Configure DDP to handle unused parameters
ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
accelerator = Accelerator( accelerator = Accelerator(
gradient_accumulation_steps=cfg.gradient_accumulation_steps, gradient_accumulation_steps=cfg.gradient_accumulation_steps,
mixed_precision=cfg.mixed_precision, mixed_precision=cfg.mixed_precision,
kwargs_handlers=[ddp_kwargs],
) )
device = accelerator.device device = accelerator.device
if accelerator.is_main_process: if accelerator.is_main_process: