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?