Maybe torch.autograd.profiler can help, but I haven’t used torch.autograd. But you can do it yourself with the time library.
For the forward pass, you can manually place time.start() and time.end(). And for the backward pass you can use hooks. Refer to this post for reference link