I implemented a custom layer and I’m having trouble getting the optimizer to work. I create my optimizer
optimizer = optim.SGD(net.parameters(), lr=0.01)

and my training loop which consists of this
optimizer.zero_grad()  
output = net(input)
loss = criterion(output, target)
loss.backward()
optimizer.step()  

doesn’t update the weights. I print the weights after initializing the model, then train, and print the weights again and they are the same. I also printed the gradients during training and they are nonzero.
However when I run
learning_rate = 0.01
for f in net.parameters():
    f.data.sub_(f.grad.data * learning_rate)

it does update the weights. Does anyone know why this is happening?