How can I force a member variable to appear as nested?

[miccio-dk]

Hi :) As per title, I have a model containing a couple of custom layers, implemented by inheriting from nn.Module. However, they do not appear as nested blocks in my graph, even tho I have "expand_nested=True". All the relevant operations performed by the module do appear, but they're just thrown around based on the execution order. How can I ensure that a given module/element appears in the graph as a nested module?

Thank you so much in advance for the amazing package (and hopefully for the help ;) )

[miccio-dk]

I figured out that the layer needs to call its forward function to be properly featured in the graph. I was calling a separate method (with a different name) and that caused it to not be recognized by torchview.