Visualization Tools that can show proper branch and merge in Graphs
You could use Tensorboard, via the tensorboardX interface.

That allows you to load models from PyTorch (and Chainer, MXNet etc.) into Tensorboard. This will then show you the full graph interactively.

From the homepage: