From f62927f71e8ad56033c7c13d613dc36c6d78993c Mon Sep 17 00:00:00 2001
From: Shu Yang <shyang@davidson.edu>
Date: Mon, 22 Jun 2026 03:42:33 +0000
Subject: [PATCH] fix(generate_wan): avoid process0-only gate for profiling run

---
 src/maxdiffusion/generate_wan.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py
index a30c4d633..7cfba422d 100644
--- a/src/maxdiffusion/generate_wan.py
+++ b/src/maxdiffusion/generate_wan.py
@@ -353,7 +353,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
   s0 = time.perf_counter()
   # Restore original profiler setting for the profiling run
   config.get_keys()["enable_profiler"] = original_enable_profiler
-  if max_utils.profiler_enabled(config):
+  if original_enable_profiler:
     # Injecting user requested XLA tracing flags
     xla_flags = os.environ.get("XLA_FLAGS", "")
     new_flags = "--xla_enable_mxu_trace=true --xla_jf_dump_llo_html=true --xla_tpu_enable_llo_profiling=true"