Hi,
Recently I’m doing some research on network pruning and built a custom layer for my vgg network.
My custom layer Mask contains a set of parameters that have the same shape as the input feature map. However, after I implemented the module together with my other conv layers using nn.Sequential, the parameters  self.mask does not get updated.
I think perhaps my custom module broke some rules in Pytorch, can you please give me some advice on getting it work?
My custom layer is defined as following:
class Mask(nn.Module):
    def __init__(self,bits=1):
        super(Mask, self).__init__()
        self.mask = None
        self.bits = bits

        self.weight_shape = None
        self.channel = None
        self.height = None
        self.width = None

    def forward(self,input):
        if type(self.mask) == type(None):
            self.mask = Parameter(torch.Tensor(input.shape[1:]).cuda())
            nn.init.constant_(self.mask,val=1.0)
            self.weight_shape = input.shape[1:]
            self.channel = input.shape[1]
            self.height = input.shape[2]
            self.width = input.shape[3]
        result = input*self.mask
        return result

the Mask layers are built using:
def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 3
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            mask = utils.Mask(bits=1).cuda()
            conv2d = utils.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [mask,conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [mask,conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

I initially defined the self.mask as None and initialized it in the forward() function because my intention was to mask the input of the layer, so I cannot fetch the size of input during initializing the module. But I suspect there might be something wrong here. If so, is there any other way I can do it correctly? Thanks in advance!