Printing a model will give you a “visualization” of how the modules and submodules are composed. Given the flexibility of the computation (which might reorder modules etc.),
this isn’t authoritative in any way, but for a well designed model (and yours would be, would it?), it goes a long way.
There also is considerably lower level way of visualizing the computational graph recorded by autograd, see e.g. (but a bit old) https://github.com/szagoruyko/pytorchviz/blob/master/torchviz/dot.py 4 (it’s a whole package with instructions how to use etc, but the linked file has the meat, really).
Best regards
Thomas