Torchviz showing layer names with nn.Module

Hey, is there a way on getting torchviz to show the layer names while using nn.Module,
All the examples are I’ve found are using nn.Sequential! for example here where the blue boxes show W0.weight, W0.bias etc.

What exactly are you trying to achieve?
Are you trying to have a sequential model where some layers are custom?
Or will the whole model not use sequential?
What might be interesting is to look at how torchviz gets the names.
What I e.g. found is the method get_var_name in
You then have to make sure that the names you want are returned by the functions (or stored in the variables) that torchviz is using.