Skip to content

Commit f62927f

Browse files
committed
fix(generate_wan): avoid process0-only gate for profiling run
1 parent 9616d1c commit f62927f

1 file changed

Lines changed: 1 addition & 1 deletion

File tree

src/maxdiffusion/generate_wan.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -353,7 +353,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
353353
s0 = time.perf_counter()
354354
# Restore original profiler setting for the profiling run
355355
config.get_keys()["enable_profiler"] = original_enable_profiler
356-
if max_utils.profiler_enabled(config):
356+
if original_enable_profiler:
357357
# Injecting user requested XLA tracing flags
358358
xla_flags = os.environ.get("XLA_FLAGS", "")
359359
new_flags = "--xla_enable_mxu_trace=true --xla_jf_dump_llo_html=true --xla_tpu_enable_llo_profiling=true"

0 commit comments

Comments
 (0)