When using net.parameters() in the optimizer, it works fine.
customlayer defn is just element-wise product, as defined below:
class CustomLayer(nn.Module):
    def __init__(self, args1, args2):
        super(CustomLayer, self).__init__()
        self.weights = nn.Parameter(torch.Tensor(args1, args2))

    def forward(self, x):
        return x * self.weights