Work on the run while it is still going.

import weightslab as wl

# <your code>: cfg, model, optimizer, train_ds ...

# wrap what you already have: one call per object
hp        = wl.watch_or_edit(cfg, flag="hyperparameters")
model     = wl.watch_or_edit(model, flag="model", device=device)
optimizer = wl.watch_or_edit(optimizer, flag="optimizer")
loss      = wl.watch_or_edit(
    nn.CrossEntropyLoss(reduction="none"),   # one value per sample
    flag="loss", signal_name="train/loss", log=True)
loader    = wl.watch_or_edit(
    train_ds, flag="data", loader_name="train_loader", is_training=True,
    batch_size=hp["data"]["train_loader"]["batch_size"], shuffle=True)

wl.serve(serving_grpc=True, serving_cli=True)

# your training loop
for epoch in range(epochs):
    for step in range(steps_per_epoch):
        with wl.guard_training_context:
            inputs, uids, labels = next(loader)
            optimizer.zero_grad()
            logits = model(inputs)
            loss_per_sample = loss(logits, labels, batch_ids=uids,
                                   preds=logits.argmax(1, keepdim=True))
            loss_per_sample.mean().backward()
            optimizer.step()
            # <your code>: logging, anything else per step
    # <your code>: evaluation, checkpoints