I want to implement a network like a fork-join model. Just like shown in the image

Untitled presentation.png960×540 40.3 KB

I want to know how to implement it on pytorch?
I have tried the following implementation but I am not getting the desired results. Is their something wrong on the implementation level.
class Network(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        # pre-trained features
        backbone = vgg16(is_caffe=True)
        
        l7 = [('fc7', nn.Conv2d(512, 1, 1))]
        l8 = [('fc8', nn.Conv2d(3, 3, 1))]
        self.encoder = copy.deepcopy(backbone)
        self.conv1_1_1=nn.Sequential(OrderedDict(l7))
        self.conv1_1_2=nn.Sequential(OrderedDict(l8))

         
    def forward(self, x, weights):
        x = self.encoder(x)
        x=x/x.max()
        #fork here
        # weights do not have any gradient flow thru them
        suppQueryFusion = [torch.add(x,weight.repeat(1, 1, x.size(2), x.size(3))) for weight in weights]
        
        weightedFeats = [self.conv1_1_1(feat) for feat in suppQueryFusion]
        # join here
        concatedFeats = torch.cat( weightedFeats,dim=1)

        x=nn.Sigmoid()(ProbabilityMap)
        x=self.conv1_1_2(x)
        return x