I’m trying to use the profiler, and I’m having trouble interpreting the output.  Is it possible to aggregate the profiling data by the python class(es) which invoked the operation?  I’d love to be able to see how much time/memory is going into MultiHeadAttention instead of just learning that “matmul” is the biggest hog.
If the profiling captures the call stack when operations are invoked I’d guess this wouldn’t be too hard.  If it’s not capturing the stack, I’d certainly be willing to pay whatever performance penalty is needed while running the profile.