Profiling Granularity Using torch profiler

Hi everyone!

I’m trying to profile different sections of my code with the torch profiler. But while I want this functionality, I also want to avoid a lot of if-else switches. I did find a way around it by actually splitting the runner loop into a runner loop that only runs for profiling and another one which runs during training. But now I wish to create a more fine grained profiling and want to document different sections as named sections. For example, just the backbone implementation separately. But there are two problems at hand: One is that the section lies outside the runner file so I need a way to pass the profiler object to the forward/backward section and another issue is that I run into a situation with too many indents and if statements within my code which render the code sections almost unreadable or too ugly. Please let me know if you guys have any suggestions or ways to actually work around this.

Please ask me any questions if it helps. Thank you!