Measure time for backward pass in each module


I am measuring time in both forward pass and backward pass in each module. For forward pass, it is easy as we can measure the time taken in the forward method.

I am wondering what I should do if I want to do the same for the backward pass?

Invoke backward() in autograd.profiler.profile context, instead of doing ad-hoc measurements. Then if you watch timeline from prof.export_chrome_trace(), you can maybe match operations with modules. I first thought that profiler.record_function() would allow tagging modules, but it seems to not work during backward.