Construct graph from Neural network

I am wondering How we could construct a graph from neural network where layer are nodes.

I have some vague idea of how to do this. For example, we can use hooks to record the layer and their connection, but this method seems to only work for sequential model, when we have skip connections, this does not work well.

I have also seen people using onnx and traces to construct a trace graph, but in that case, the node are onnx operation instead of pytorch module.