Solved by ptrblck in post #2 
              
              
                Your custom layer contains the weigth parameter printed by this code snippet: 
class CustomLayer(torch.nn.Module):
    def __init__(self, input_shape):
        super(CustomLayer,self).__init__()
        zeros = torch.zeros(input_shape)
        self.weights = torch.nn.Parameter(zeros)
        torch.n…