If you want to pass the mask as an additional argument to the forward method of your custom Dropout layer, you should write this logic into the forward method of VGG explicitly and remove the usage of the nn.Sequential block as it expects a single input/output in the default implementation.