Solved by vmirly1 in post #2 
              
              
                When you define a new custom layer in a class, you can inherit the properties of nn.Module and define the learnable parameters as nn.Parameter(). Then, these parameters can be updated by the autograd module. I have written a small example for a custom defined layer: 
import torch
import torch.nn as …