"""Train a tiny torch model with launchhelm.torch.

Requires ``pip install "launchhelm[torch]"`` and the same LAUNCHHELM_*
environment variables as log_run.py.

    python torch_quickstart.py
"""

import torch

import launchhelm.torch as lht

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Linear(8, 16), torch.nn.ReLU(), torch.nn.Linear(16, 1)
)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
x = torch.randn(64, 8)
y = x.sum(dim=1, keepdim=True)

tr = lht.init("quickstart", name="torch-quickstart", config={"lr": 0.01})
tr.watch(model, optimizer)
dropout_p = tr.tunable("dropout_p", 0.0, minimum=0.0, maximum=0.9)

for step in range(20):
    for module in model.modules():
        if isinstance(module, torch.nn.Dropout):
            module.p = dropout_p.value
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    tr.clip_grad_norm_()
    optimizer.step()
    tr.step(step=step, loss=loss.item())

tr.finish()
print("run finished:", tr.run.run_id)