I need to keep the gradient values of my model (a convolutional neural network). The problem is when I use the autograd.grad function, the optimizer.step() doesn't seem to work (the loss function doesn't change after the optimization step). So my question is how to use the autograd.grad function with the optimizer.step()?
def train(epoch):
optimizer.zero_grad()
out = net(gt_data)
y = criterion(out, gt_onehot_label)
train_losses.append(y)
y.backward(retain_graph=True)
dy_dx = torch.autograd.grad(y, net.parameters())
optimizer.step()