I’ve been running over this problem for the first time a few days ago and I have been unable since then to find a fix, so here I am. I am a big fan of tensorboard’s graph visualization tool, as it helps a great deal when it comes to understanding what happens in a model. I have been successfully using it for months now, but it recently tarted behaving strangely (probably since 1.2 where torch.utils.tensorboard._pytorch_graph.graph was significantly modified). My problem is that it longer nests modules inside one another. It just creates a big module, and inside are all the leaf modules directly, which is quite difficult to understand. See below for an example.
Is there any way to come back to the old behavior with the new code structure ? Because this way of having things displayed is not very handy to use to be honest, and I can’t find a way to make it work.