import torch
from torch import nn


torch.manual_seed(11)

features = torch.tensor(
    [
        [0.0, 0.5],
        [1.0, -0.5],
        [2.0, 1.0],
    ],
    dtype=torch.float32,
)
targets = torch.tensor([[0.2], [0.7], [1.6]], dtype=torch.float32)

model = nn.Linear(2, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)


def grad_summary(label):
    pieces = []
    for name, parameter in model.named_parameters():
        if parameter.grad is None:
            pieces.append(f"{name}=None")
        else:
            pieces.append(
                f"{name}_norm={parameter.grad.norm().item():.4f}; "
                f"{name}_sum={parameter.grad.sum().item():.4f}"
            )
    print(f"{label}: {' | '.join(pieces)}")


grad_summary("before backward")

loss = loss_fn(model(features), targets)
loss.backward()
grad_summary("after backward")

optimizer.step()
optimizer.zero_grad()
grad_summary("after zero_grad default")

loss = loss_fn(model(features), targets)
loss.backward()
grad_summary("after next backward")

optimizer.zero_grad(set_to_none=False)
grad_summary("after zero_grad set_to_none_false")