From bb824f2275683b9122e44a7f6815e84145c02eef Mon Sep 17 00:00:00 2001 From: Pepijn Date: Tue, 14 Oct 2025 14:08:20 +0200 Subject: [PATCH] change accelerate detection --- src/lerobot/utils/utils.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/lerobot/utils/utils.py b/src/lerobot/utils/utils.py index b32decbb9..1cdfed67c 100644 --- a/src/lerobot/utils/utils.py +++ b/src/lerobot/utils/utils.py @@ -179,7 +179,19 @@ def format_big_number(num, precision=0): def is_launched_with_accelerate() -> bool: - return "ACCELERATE_MIXED_PRECISION" in os.environ + """Check if the script was launched using accelerate. + + Accelerate sets several environment variables when launching distributed training. + We check for variables that are always set regardless of the configuration. + """ + # Check for environment variables that accelerate always sets + accelerate_env_vars = [ + "ACCELERATE_MIXED_PRECISION", # Set when mixed precision is configured + "LOCAL_RANK", # Always set in distributed training + "WORLD_SIZE", # Always set in distributed training + "ACCELERATE_USE_CPU", # Set by accelerate + ] + return any(var in os.environ for var in accelerate_env_vars) def say(text: str, blocking: bool = False):