I have defined model using nn.module. I need to use grads() function instead of backward(). But unfortunaltely I am not getting the right results. The model is not converging. Here is snippet of my code for SGD implemenation. Is there any problem here?
grads = grad(loss, model.parameters())
with torch.no_grad():
for (W, g) in zip(model.parameters(), grads):
W.data -= step_size*g