Layer wise backward pass profiling

I’m trying to figure out a way to do layer wise backward pass profiling.
I’ve used register_backward_hook(hook_func)
and have been trying to use NVTX push and pop in the hook_func.

however, since the hook_func is executed after the backward pass for a certain layer,
it seems that it is impossible to properly annotate each layer for backward pass.

is there another way to profile layer wise backward pass using NVTX?