I have four parameters ‘a’, ‘b’, ‘weight’, and, ‘bias’ defined in the PyTorch module.
In the forward pass,
Step 1:    Creating a sparse matrix ‘T’  that is dependent on the parameters ‘a’ and ‘b’
Step 2:    mat_1 = torch.mm (input, weight)
Step 3:    output = SparseMM()(T, mat_1)
Step 4:    return output + bias
where, SparseMM is an autograd function implemented as below:

class SparseMM(torch.autograd.Function):
“”"
Sparse x dense matrix multiplication with autograd support.
Implementation by Soumith Chintala:
https://discuss.pytorch.org/t/
does-pytorch-support-autograd-on-sparse-matrix/6156/7
"""

def forward(self, matrix1, matrix2):
    self.save_for_backward(matrix1, matrix2)
    return torch.mm(matrix1, matrix2)

def backward(self, grad_output):
    matrix1, matrix2 = self.saved_tensors
    grad_matrix1 = grad_matrix2 = None

   
    if self.needs_input_grad[1]:
        grad_matrix2 = torch.mm(matrix1.t(), grad_output)

    return grad_matrix2


The weight and bias are the parameters that get updated, but the parameter ‘a’ and ‘b’ that create the sparse matrix doesn’t get updated.
Is there a way to update these parameters as well?