import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


torch.manual_seed(23)

features = torch.tensor(
    [
        [-1.0, 0.0, 0.5],
        [-0.5, 0.25, 1.0],
        [0.0, -0.5, 0.25],
        [0.5, 0.75, -0.25],
        [1.0, -0.25, -0.5],
        [1.5, 0.5, 0.0],
        [2.0, -0.75, 0.75],
        [2.5, 1.0, -1.0],
    ],
    dtype=torch.float32,
)
targets = features @ torch.tensor([[0.8], [-0.4], [0.3]]) + 0.2

train_data = TensorDataset(features, targets)
train_loader = DataLoader(train_data, batch_size=4, shuffle=False)

model = nn.Sequential(nn.Linear(3, 1))
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.08)
initial_weight = model[0].weight.detach().clone()