Visualization Tools that can show proper branch and merge in Graphs

I want help finding the visualization tool can draw similar architecture as given in the image below

Keras visualization produces something similar to this graph. But I'm working in Pytorch. I tried using Netron and Graphviz, both produce graphs that do not show the branching and merging properly.

This is the plot I rendered in Graphviz for another branching and merging architecture.

Here's the python code I used to render the plot

graph {
    graph [rankdir=RL]
    input [label=Input]
    conv1 [label="Conv 1"]
    conv2 [label="Conv 2"]
    conv3 [label="Conv 3"]
    conv4 [label="Conv 4"]
    conv5 [label="Conv 5"]
    conv6 [label="Conv 6"]
    conv7 [label="Conv 7"]
    conv8 [label="Conv 8"]
    conv9 [label="Conv 9"]
    batch1 [label="Batch Norm1"]
    batch2 [label="Batch Norm2"]
    batch3 [label="Batch Norm3"]
    caps1_a [label="Caps 1(a)"]
    caps1_b [label="Caps 1(b)"]
    caps2_a [label="Caps 2(a)"]
    caps2_b [label="Caps 2(b)"]
    caps3_a [label="Caps 3(a)"]
    caps3_b [label="Caps 3(b)"]
    sum1 [label=Sum]
    sum2 [label=Sum]
    sum3 [label=Sum]
    stack [label=Stack]
    sum4 [label=Sum]
    softmax [label=Softmax]
    input -- conv1 [constraint=true]
    conv1 -- conv2 [constraint=false]
    conv2 -- conv3 [constraint=false]
    conv3 -- conv4 [constraint=false]
    conv4 -- conv5 [constraint=false]
    conv5 -- conv6 [constraint=false]
    conv6 -- conv7 [constraint=false]
    conv7 -- conv8 [constraint=false]
    conv8 -- conv9 [constraint=false]
    conv3 -- batch1 [constraint=false]
    conv6 -- batch2 [constraint=false]
    conv9 -- batch3 [constraint=false]
    batch1 -- caps1_a [constraint=false]
    batch2 -- caps2_a [constraint=false]
    batch3 -- caps3_a [constraint=false]
    caps1_a -- caps1_b [constraint=false]
    caps2_a -- caps2_b [constraint=false]
    caps3_a -- caps3_b [constraint=false]
    caps1_b -- sum1 [constraint=false]
    caps2_b -- sum2 [constraint=false]
    caps3_b -- sum3 [constraint=false]
    sum1 -- stack [label=w1 constraint=false]
    sum2 -- stack [label=w2 constraint=false]
    sum3 -- stack [label=w3 constraint=false]
    stack -- sum4 [constraint=false]
    sum4 -- softmax [constraint=false]
    input -- conv2 [label=k1 constraint=false]
    conv2 -- conv4 [label=k2 constraint=false]
    conv4 -- conv6 [label=k3 constraint=false]
    conv6 -- conv8 [label=k4 constraint=false]
}