Set record_shapes=True for profiler

ghstack-source-id: 6f1ed49d15ce311f1bf118820965cdb5309a8030
Pull Request resolved: #419

Author: Andrew Gu

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