import torch
from torch import nn


torch.manual_seed(7)

features = torch.tensor(
    [
        [-1.0, 0.5, 0.25],
        [0.0, -0.25, 0.75],
        [0.5, 0.75, -0.5],
        [1.0, -0.5, -0.25],
    ],
    dtype=torch.float32,
)
targets = features @ torch.tensor([[0.6], [-0.4], [0.2]]) + 0.1

model = nn.Linear(3, 1)
loss_fn = nn.MSELoss()

optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
weight_before = model.weight.detach().clone()

optimizer.zero_grad(set_to_none=True)
predictions = model(features)
loss_before = loss_fn(predictions, targets)
loss_before.backward()
optimizer.step()

with torch.no_grad():
    loss_after = loss_fn(model(features), targets)

state = optimizer.state[model.weight]
weight_changed = not torch.equal(weight_before, model.weight.detach())

print(f"optimizer={optimizer.__class__.__name__}")
print(f"learning_rate={optimizer.param_groups[0]['lr']}")
print(f"state_step={int(state['step'].item())}")
print(f"weight_changed={weight_changed}")
print(f"loss_before={loss_before.item():.6f}")
print(f"loss_after={loss_after.item():.6f}")
print(f"loss_decreased={loss_after.item() < loss_before.item()}")