poor prediction in x**2 regression in pytorch

I am trying to replicate the code available in https://machinelearningmastery.com/neural-networks-are-function-approximators/ using pytorch and adding train and test data but the prediction and loss results are not good.

In first attempt, I tried to change deep network nodes, epochs and learning rate, then
I tried to add scheduler to adjust learning rate and I implemented a small checkpoint to save the good models but still it was not enough to get good results.

I was wondering if the community have any idea to fix my code.

DATA

X = torch.arange(start, end, step, dtype=torch.float32).unsqueeze(dim=1)
y = torch.tensor([i**2.0 for i in X[0:]]).unsqueeze(dim=1)

train_split = int (0.8 *len(X))
X_train, y_train = X[:train_split], y[:train_split]
X_test , y_test = X[train_split:], y[train_split:]

scaler_x = MinMaxScaler()
scaler_x.fit(X_train)
X_Train = scaler_x.transform(X_train)


scaler_y = MinMaxScaler()
scaler_y.fit(y_train)
y_Train = scaler_y.transform(y_train)

Deep Network

class FunctionEstimatorModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.linear_layer_1 = nn.Linear(in_features = 1,out_features = 200)
        self.relu = nn.LeakyReLU()
        self.linear_layer_2 = nn.Linear(in_features = 200,out_features = 200)
        self.relu = nn.LeakyReLU()
        self.linear_layer_3 = nn.Linear(in_features = 200,out_features = 1)



    def forward (self, x: torch.Tensor) -> torch.Tensor:
        return self.linear_layer_3(self.relu(self.linear_layer_2(self.relu(self.linear_layer_1(x)))))

for loop

for epoch in range(epochs):

    model_0.train()
    y_preds = model_0(X_train).squeeze()
    loss = loss_fn(y_preds, y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    scheduler.step()


#####Testing 
    model_0.eval()

    with torch.inference_mode():
        test_pred = model_0(X_test).squeeze()
        test_loss = loss_fn(test_pred, y_test)

    print(f"Epoch: {epoch} | Loss: {loss} | Test Loss: {test_loss}") 

enter image description here

Thanks a lot

For feedback or comments, reach us on hello@newswire.ae