Dear Experts
I try to generate a simple custom linear layer as follows, but the prediction of the network is incorrect 
I tried hard for more than 2 weeks but I could not solve it. I hope someone help me.
class linearZ(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input, weight):

        ctx.save_for_backward(input, weight)
        l2IN = input
        l2 = l2IN * weight
        return l2

    @staticmethod
    def backward(ctx, grad_output):

        learnign_rate = 0.01
        input, weight = ctx.saved_tensors
        grad_weight = input * grad_output
        weight = weight - learnign_rate * grad_weight
        net.linearZ.weight.data = weight
        return grad_weight, None


class MyLinearZ(nn.Module):
    def __init__(self):
        super(MyLinearZ, self).__init__()
        self.fn = linearZ.apply
        self.weight = nn.Parameter(torch.Tensor([[np.random.randn(1, 1) * 5.66]]))

    def forward(self, x):
        x = self.fn(x, self.weight)
        return x


class Net(nn.Module):
    def __init__(self, conv_weight):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 1, 2)
        self.pool = nn.MaxPool2d((1, 3))
        self.linearZ = MyLinearZ()
        self.conv1.weight = nn.Parameter(conv_weight)



    def forward(self, x):
        x = torch.atan(self.conv1(x))
        x = self.pool(x)
        x = torch.atan(self.linearZ(x))
        return x

Thanks before all