Solved by chenglu in post #4 
              
              
                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 * inpu…