I’m not sure if you will train the share weight twice in one forward action, if it is , I think the forward code should be something like following.
def forward(self, x):
    input_image = x
    ...
    x = self.share_weight_layer(x)
    ...
    pool = x = self.avg_pool(x)
    product = pool * input_image
    x = self.some_layer(product)
    ...
    x = self.share_weight_layer(x) # this is the same layer as previous one
    ...
    return x