I use “torch.autograd.Function”  to construct my own autograd function. “eta” is my weights. But when I print all the parameters, I found eta is not updated at all.  And if I did
print(list(model.parameters())[0].grad)

I will get None.
The sample code is below:
class ANN_Pytorch(torch.nn.Module):
    def __init__(self, D_in, H1, H2, D_out, cutoff, beta):
        super(ANN_Pytorch, self).__init__()
        self.w0 = Variable(((15 - 2.5) * torch.rand(1, D_in, device=device, requires_grad=True) + 2.5)).double().unsqueeze(0).permute(2, 0, 1)
        self.eta = nn.Parameter(self.w0)
        self.linear1 = torch.nn.Linear(D_in, H1)
        self.linear2 = torch.nn.Linear(H1, H2)
        self.linear3 = torch.nn.Linear(H2, D_out)


    def forward(self, X):
        output = torch.zeros(X.shape[0], X.shape[1])
        for i in range(X.shape[0]):
            output[i,:] = RigidityExp.apply(sparse2tensor(X[i,:].reshape(1,-1)).double(), self.eta, beta, cutoff)
        z1 = self.linear1(output)
        f1 = F.sigmoid(z1)
        z2 = self.linear2(f1)
        f2 = F.sigmoid(z2)
        y_hat = self.linear3(f2)
        return y_hat
#================================Construct Model================================
D_in, H1, H2, D_out, cutoff, beta = X_train_Distance.shape[1], 40, 20, ytrain.shape[1], 12, 2.5

model = ANN_Pytorch(D_in, H1, H2, D_out, cutoff, beta).cuda()

criterion = torch.nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr = 1e-2, momentum = 0.9, weight_decay = 1e-2)
epochs = 300

#===================================Trianing Model==============================
for epoch in range(epochs):
    model.train()
    y_pred  = model(X_train_Distance)
    loss = criterion(y_pred, ytrain)
    print(epoch, loss.item())
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(list(model.parameters())[0].grad)