From 0852287c73f85a35a579f5a9ad09feed316d2d14 Mon Sep 17 00:00:00 2001
From: Andrew Gu <andgu@fb.com>
Date: Mon, 24 Jun 2024 13:44:33 -0700
Subject: [PATCH] Set `record_shapes=True` for profiler

[ghstack-poisoned]
---
 torchtitan/profiling.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/torchtitan/profiling.py b/torchtitan/profiling.py
index c993a74f5..662b64f8c 100644
--- a/torchtitan/profiling.py
+++ b/torchtitan/profiling.py
@@ -68,6 +68,7 @@ def trace_handler(prof):
             ],
             schedule=torch.profiler.schedule(wait=wait, warmup=warmup, active=active),
             on_trace_ready=trace_handler,
+            record_shapes=True,
         ) as torch_profiler:
             torch_profiler.step_num = global_step
             yield torch_profiler