합치기 — Production-Grade Loop

import math
import torch
import torch.nn as nn
from torch.amp import autocast
from tqdm import tqdm

# 1. Device
device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)

# 2. Model
model = MyModel().to(device)

# 3. Loss
criterion = nn.CrossEntropyLoss()

# 4. Optimizer with weight-decay split
decay, no_decay = [], []
for n, p in model.named_parameters():
    if not p.requires_grad: continue
    if p.dim() < 2 or any(k in n for k in ('bias', 'norm')):
        no_decay.append(p)
    else:
        decay.append(p)
optimizer = torch.optim.AdamW(
    [{'params': decay, 'weight_decay': 0.01},
     {'params': no_decay, 'weight_decay': 0.0}],
    lr=1e-4,
)

# 5. Scheduler — warmup + cosine
total_steps = len(train_loader) * num_epochs
warmup_steps = total_steps // 20    # 5% warmup
def lr_lambda(step):
    if step < warmup_steps:
        return step / warmup_steps
    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    return 0.5 * (1.0 + math.cos(math.pi * progress))
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

best_val = float('inf')
global_step = 0

for epoch in range(num_epochs):
    # ---- TRAIN ----
    model.train()
    pbar = tqdm(train_loader, desc=f"epoch {epoch:02d}")
    running = 0.0

    for batch_x, batch_y in pbar:
        batch_x = batch_x.to(device, non_blocking=True)
        batch_y = batch_y.to(device, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)

        with autocast(device_type=device, dtype=torch.bfloat16) if device != 'cpu' else nullcontext():
            output = model(batch_x)
            loss = criterion(output, batch_y)

        if not torch.isfinite(loss):
            print(f"step {global_step}: non-finite loss; skipping")
            continue

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        scheduler.step()
        global_step += 1

        running = 0.95 * running + 0.05 * loss.item() if global_step > 1 else loss.item()
        pbar.set_postfix(loss=f"{running:.4f}", lr=f"{scheduler.get_last_lr()[0]:.2e}")

    # ---- VALIDATE ----
    model.eval()
    val_loss = 0.0
    val_n = 0
    with torch.inference_mode():
        for vx, vy in val_loader:
            vx, vy = vx.to(device), vy.to(device)
            val_loss += criterion(model(vx), vy).item() * vx.size(0)
            val_n += vx.size(0)
    val_loss /= val_n
    print(f"epoch {epoch:02d}  val_loss={val_loss:.4f}")

    # ---- CHECKPOINT ----
    if val_loss < best_val:
        best_val = val_loss
        torch.save({
            'epoch': epoch,
            'val_loss': val_loss,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict(),
        }, 'best.pt')

from contextlib import nullcontext

# autocast doesn't apply on CPU. Use nullcontext as a no-op stand-in
# so the same training code runs cleanly on CPU for debugging.

device = "cpu"  # for example
amp_ctx = autocast(device_type=device, dtype=torch.bfloat16) if device != 'cpu' else nullcontext()

with amp_ctx:
    out = model(x)
    loss = criterion(out, y)