import torch
from torch import nn


torch.manual_seed(11)

features = torch.tensor(
    [
        [0.0, 0.5],
        [1.0, 1.5],
        [2.0, 2.5],
        [3.0, 3.5],
    ],
    dtype=torch.float32,
)
weights = torch.tensor([[0.4], [-0.2]])
targets = features @ weights + 0.1

model = nn.Linear(2, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
scheduler = torch.optim.lr_scheduler.StepLR(
    optimizer,
    step_size=2,
    gamma=0.5,
)

print(f"initial lr: {scheduler.get_last_lr()[0]:.4f}")

observed_lrs = []

for epoch in range(1, 6):
    optimizer.zero_grad(set_to_none=True)
    loss = loss_fn(model(features), targets)
    loss.backward()
    optimizer.step()
    scheduler.step()
    current_lr = scheduler.get_last_lr()[0]
    observed_lrs.append(round(current_lr, 4))

    print(
        f"epoch {epoch}: "
        f"loss={loss.item():.4f}, "
        f"lr={current_lr:.4f}"
    )

expected_lrs = [0.05, 0.025, 0.025, 0.0125, 0.0125]
matched = observed_lrs == expected_lrs

print(f"expected: {expected_lrs}")
print(f"observed: {observed_lrs}")
print(f"matched: {matched}")