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.Sequential(nn.Linear(3, 1))
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

first_weight = model[0].weight
weight_before = first_weight.detach().clone()

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

gradient_norm = torch.linalg.vector_norm(first_weight.grad).item()
optimizer.step()

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

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

print(f"torch_version={torch.__version__}")
print(f"optimizer={optimizer.__class__.__name__}")
print(f"learning_rate={optimizer.param_groups[0]['lr']}")
print(f"loss_before={loss_before.item():.6f}")
print(f"gradient_norm={gradient_norm:.6f}")
print(f"state_step={int(state['step'].item())}")
print(f"exp_avg_norm={torch.linalg.vector_norm(state['exp_avg']).item():.6f}")
print(f"weight_changed={weight_changed}")
print(f"loss_after={loss_after.item():.6f}")
print(f"loss_decreased={loss_after.item() < loss_before.item()}")