Decorator for profiling
How about having a decorator suitable for any function / class method that measures the time each function call takes and if a logger provided, logs it (for example to tensorboard). If you then just split up you functionality in your code in "meaningful" snippets (functions / methods) you would just have to append that decorator to each part you want to profile.
It would also be nice to have the possibility to "turn off" the logging globally so that you would not have to remove all the decorators if you are finished debugging / profiling.