Solved by ptrblck in post #2 
              
              
                You could iterate the modules of the parent nn.Module, check for the mask attribute via hasattr and assign e.g. a new nn.Parameter to it as seen here: 
class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.mask = None
        

class MyModel(nn.Module):
    def _…