I use this code (but tweaked a little): https://github.com/szagoruyko/functional-zoo/blob/master/visualize.py
Since PyTorch is a dynamic framework there isn’t really a graph like in TensorFlow / Keras. There is only the graph that was created when you did some computation. In the case of a neural network, that is the computations for when you did a forward pass.
The make_dot()
function from that source code takes the output of your NN (such as the loss) and then draws the graph that was used to compute that loss. However, it only knows about the low-level PyTorch operations, so you may get more details than you want.