Thanks for the tip.  I’m trying to understand it, as the docs are pretty thin.  Would I do something like this:
    torch.cuda.nvtx.range_push("multihead_attention")
    out = self.attention(x)
    torch.cuda.nvtx.range_pop("multihead_attention")

And then hopefully this shows up in the profiler output?