Thanks @tom,
I really appreciate the detailed answer!
I believe that this is important to support, to allow plotting aiding in having a full mental image of what’s going on in the network that you are training (including parameters, buffers and activations).
Depending on ONNX or JIT is problematic because I don’t think it can support all of the exciting custom stuff that is so easy to do in pytorch (correct me if I’m wrong).
I would love to add the implementation, but it sounds like a fairly delicate topic to jump into, without having any experience in autograd CPP code.
However, with some guidance I can give it a try 
Any other suggestions on other approaches to plot a graph that contains operators, parameters, buffers and activations ?
Maybe I should give the JIT tracing more chance first ?