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