I’m really unsure so I’ll just post some ideas.
Since the computation graph is created during the forward pass, it might of course be different for each pass (e.g. if you are using conditions etc.).
You could try to use the grad_fn and call into grad_fn.next_functions to crawl the graph.
However, this would yield you the operations, not necessarily the layers and I’m not sure how hard it would be to create the mapping. 