Pytorch Does Not Converge When Approximating Square Function With Linear Model

I'm trying to learn some PyTorch and am referencing this discussion here The author provides a minimum working piece of code that illustrates how you can use PyTorch to solve for a

Solution 1:

You cannot fit a 2nd degree polynomial with a linear function. You cannot expect more than random (since you have random samples from the polynomial).
What you can do is try and have two inputs, x and x^2 and fit from them:

model = nn.Linear(2, 1)  # you have 2 inputs now
X_input = torch.cat((X, X**2), dim=1)  # have 2 inputs per entry
# ...

    predictions = model(X_input)  # 2 inputs -> 1 output
    loss = loss_fn(predictions, t)
    # ...
    # learning t = c*x^2 + a*x + b
    print("learned a = {}".format(list(model.parameters())[0].data[0, 0]))
    print("learned c = {}".format(list(model.parameters())[0].data[0, 1])) 
    print("learned b = {}".format(list(model.parameters())[1].data[0]))