I am building a really simple model to learn the parameter in a Poisson model and I am not sure where I am going wrong. I am using pytorch.nn and doing the following.
I made some really simple fake data
# This is the value I am trying to estimate

x = torch.tensor(2.0)


# This is a value drawn from the Poisson(x) distribution 
# In this example it is 4

y = torch.poisson(x).reshape(1)

Then I just set up a really simple model
# I initialised the parameter that is going to estimate x with a random value (0.2) 
# and set that it requires a gradient

a = torch.tensor([0.2], requires_grad = True)


# I define the loss function with log_input set to false

loss_function = torch.nn.PoissonNLLLoss(log_input = False)


# Defined the model

def model(a):
    return torch.poisson(a)

# And the parameter to be optimised 
# I chose SGD arbitrarily, maybe this is the problem?

optimizer = torch.optim.SGD([a], lr = 0.1)

Then I do iterations to update a
for i in range(2000):
    
    # Forward pass

    y_pred = model(a)


    # Compute the loss

    loss = loss_function(y_pred, y)
    

    # Backprop

    optimizer.zero_grad()
    
    loss.backward()
    
    
    # Update parameters

    optimizer.step()

The problem is after this the a is still 0.2 and if I call a.grad it is 0. Where am I going wrong?
I have also tried instead to initiate a class for the model inheriting a nn.Module. However the same problem persists :
class learning_model(nn.Module):
    
    def __init__(self):
        super().__init__()
        self.a = nn.Parameter(torch.rand(1))
        self.a.requires_grad = True
        
    def forward(self):
        return torch.poisson(self.a)

model = learning_model()

loss_function = nn.PoissonNLLLoss(log_input = False)

optimizer = torch.optim.SGD(model.parameters(), lr = 0.1)

print(model.a)

Outputs:
Parameter containing:
tensor([0.1402], requires_grad=True)

Then:
for i in range(20):
    
    # Forward pass
    y_pred = model()
    
    # Compute the loss
    loss = loss_function(y_pred, y)
    
    # Backprop
    optimizer.zero_grad()
    
    loss.backward()    
    
    # Update parameters
    optimizer.step()
    
print(model.a, '\n gradient:', model.a.grad)

Outputs:
Parameter containing:
tensor([0.1402], requires_grad=True) 
 gradient: tensor([0.])