Neural network not training, parameters.grad is None

Question

I implemented a simple NN and a custom objective function for a minimization problem. Everything seems to be working fine, except for the fact that the model does not seem to learn.

I checked if list(network.parameters())[0].grad was None, and indeed this seems to be the problem. Based on previously asked questions, the problem seems to be the detachment of the graph, but I don't know what I am doing wrong.

Here's the link to the code that you can run on colab: Colab code

Thank you!!

Answer 1

This part seems problematic in your code.

  output_grad, _ = torch.autograd.grad(q, (x,y))
  output_grad.requires_grad_()

Your loss depends on output_grad and so when you do loss.backward() are trying to compute the gradient of parameters w.r.t to output_grad.

You cannot compute the gradient of output_grad since create_graph is False by default. And so output_grad is implicitly detached from the rest of the graph. To fix this, just pass create_graph=True in the autograd.grad. You do not need to set requires_grad either for output_grad, i.e., the second line is not needed.