xla

d8125517 - [FSDP] Add profiler to ImageNet (#4754)

[FSDP] Add profiler to ImageNet (#4754) Summary: This pull request adds profiler to test/test_train_mp_imagenet_fsdp.py, and moves all the tracing part into the build_graph closure in test_train_mp_imagenet.py. Test Plan: CI.