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