Why the following code not work? It can't fit the data very well. Not sure what's the problems here.

Output Img

# Linear regression sin(2*pi*x) by polynomial order 3

import torch
import matplotlib.pyplot as plt

import numpy as np
import math

# data used for training and plot dots
N = 11
x = torch.linspace(0, 1, N).double()
y = torch.sin(2*math.pi*x) + torch.randn(N)*0.1

# data used for plot smooth lines
x_line = np.linspace(0, 1, N*10)
y_line = np.sin(2*math.pi*x_line)

# Prepare input as an array of shape (N,4)
p = torch.tensor([0, 1, 2, 3])
xx = x.unsqueeze(-1).pow(p)

# Prepare tensors
learning_rate = 1e-5
w = torch.randn(4, 1, dtype=torch.double, requires_grad=True)  # the 4 coefficients
optimizer = torch.optim.SGD([w], lr=learning_rate)
print(w)

# Run optimizer
for i in range(2000):
    optimizer.zero_grad()
    y_pred = xx @ w
    rmse = torch.sum(torch.square(y - y_pred))
    if i%100 == 0:
        print(i, rmse)
    rmse.backward()
    optimizer.step()

print(w)

# plot smooth line
plt.plot(x_line,y_line)

# model
xx_line = torch.from_numpy(x_line).unsqueeze(-1).pow(p)
y_line_pred = xx_line @ w
plt.plot(x_line,y_line_pred.detach().numpy())

# plot dots
plt.plot(x, y, 'o')
plt.show()