[diffusion] fix: Fix profiler trace missing Python stack in diffusion pipeline (#14499)

This commit is contained in:
Xiaoyu Zhang
2025-12-05 20:12:35 +08:00
committed by GitHub
parent 2ce121a1c3
commit 5347732219

View File

@@ -599,9 +599,7 @@ class DenoisingStage(PipelineStage):
active=batch.num_profiled_timesteps,
repeat=5,
),
on_trace_ready=lambda _: torch.profiler.tensorboard_trace_handler(
f"./logs"
),
on_trace_ready=None,
record_shapes=True,
with_stack=True,
)