Hi, I am trying to learn a scalar parameter to define a matrix, so I have written a custom autograd function that allows me to do the type casting, since that is non-differentiable. Autograd calls the backward function, so there is a gradient, but the parameter itself doesn’t update (tried different learning rates already).
class ExampleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.prob = torch.nn.Parameter(torch.tensor(0.5, dtype=torch.float32), requires_grad=True)

    def forward(self, x):
        output = quantize_length(self.prob).cuda()
        return output


def linear_length(prob):
    # [0, 8] length in frames = 9 possible options
    # input within range [0 1], map to an integer length
    return torch.floor(9 * prob).type(torch.uint8)


class QuantizeLengthFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, length_prob):
        ctx.save_for_backward(length_prob)
        output = linear_length(length_prob)  # quantize length
        weight = torch.zeros((10, 10), dtype=torch.float32)
        weight[:, :output] = torch.ones((10, output), dtype=torch.float32).cuda()
        return weight

    def backward(ctx, grad_output):
        length_prob = ctx.saved_tensors
        grad_input = grad_output.clone().sum() * 9.0
        return grad_input

quantize_length = QuantizeLengthFunction.apply

# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 64, 15, 10, 10

# Create random Tensors to hold input and outputs.
x = torch.rand(N, D_in, device=device, dtype=dtype) - 0.5
y = torch.rand(N, H, device=device, dtype=dtype) - 0.5

# Create random Tensors for weights.
w1 = torch.rand(D_in, H, device=device, dtype=dtype, requires_grad=True)

model = ExampleModel()

for name, data in model.named_parameters():
    print(name)
optim = torch.optim.Adam(model.parameters(), lr=1e-4, eps=1e-8)

for t in range(500):

    output = model(x)
    y_pred = (x.mm(w1)).mm(output)
    loss = (y_pred - y).pow(2).sum()

    # Use autograd to compute the backward pass.
    optim.zero_grad()
    loss.backward()
    optim.step()

If instead of using autograd to compute the backward pass, I use the following then the update works:
    # # Update weights using gradient descent
     with torch.no_grad():
         w1 -= learning_rate * w1.grad
         prob -= learning_rate * prob.grad
    
         # Manually zero the gradients after updating weights
         w1.grad.zero_()
         prob.grad.zero_()

But I’m trying to put this in a much larger network so I’d like to use Autograd to update this scalar along with the rest of my network. What am I doing wrong here?