Function train_loop

examples/cpu/cpu_training.py:114–145

(model, optimizer, dataloader, steps, log_interval)

def train_loop(model, optimizer, dataloader, steps, log_interval):
    model.train()
    history = []
    step = 0
    t0 = time.time()

    while step < steps:
        for batch in dataloader:
            if step >= steps:
                break

            input_ids = batch["input_ids"]
            attention_mask = batch["attention_mask"]
            labels = batch["labels"]

            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss
            loss.backward()

            optimizer.step()
            optimizer.zero_grad()

            loss_val = loss.item()
            elapsed = time.time() - t0
            history.append((step, loss_val, elapsed))

            if step % log_interval == 0:
                print(f"  step {step:4d} | loss {loss_val:.4f} | time {elapsed:.1f}s")

            step += 1

    return history


def get_torch_dtype(name):