Solved by beaupreda in post #2 
              
              
                Hello, 
Here’s what I would do (assuming I understood your problem correctly): 
class NoiseLayer(nn.Module):
    # I do not think you need the x and y dimensions here
    def __init__(self, noise_percentage):
        super(NoiseLayer, self).__init__()
        self.n_p = noise_percentage

    def for…