suggestion of adding torch.profile

i just check out the code of appendix-A/01_main-chapter-code /DDP-script.py,how about adding 
```
from torch.profiler import profile
with profile() as prof:
    #the main function training code
if rank == 0:
    print("exporting trace")
    prof.export_chrome_trace("trace_ddp_simple.json")
```
than we can see the tracing profiling json file in google Chrome