import torch
from torch import nn


torch.manual_seed(23)

model = nn.Sequential(
    nn.Linear(4, 6),
    nn.ReLU(),
    nn.Linear(6, 2),
)

model[0].requires_grad_(False)

trainable_parameters = [
    (name, parameter)
    for name, parameter in model.named_parameters()
    if parameter.requires_grad
]
optimizer = torch.optim.SGD(
    [parameter for _, parameter in trainable_parameters],
    lr=0.1,
)
loss_fn = nn.CrossEntropyLoss()

features = torch.tensor(
    [
        [0.5, -1.0, 0.3, 2.0],
        [1.0, 0.2, -0.4, 0.7],
        [-0.3, 1.2, 0.8, -1.1],
        [1.5, -0.7, 0.1, 0.4],
    ],
    dtype=torch.float32,
)
targets = torch.tensor([0, 1, 1, 0])

before_step = {
    name: parameter.detach().clone()
    for name, parameter in model.named_parameters()
}

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

print(f"torch_version={torch.__version__}")
print("requires_grad:")
for name, parameter in model.named_parameters():
    print(f"  {name}={parameter.requires_grad}")

print(
    "optimizer_params="
    + ",".join(name for name, _ in trainable_parameters)
)
print(f"loss={loss.item():.4f}")

print("gradient_state:")
for name, parameter in model.named_parameters():
    if parameter.grad is None:
        print(f"  {name}=None")
    else:
        print(f"  {name}_norm={parameter.grad.norm().item():.4f}")

optimizer.step()

print("changed_after_step:")
for name, parameter in model.named_parameters():
    changed = not torch.equal(before_step[name], parameter.detach())
    print(f"  {name}={changed}")