How do you debug a PyTorch training loop where the loss is not decreasing or is NaN?

import torch
import torch.nn as nn

model     = nn.Linear(10, 5)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

for step, (X, y) in enumerate(loader):
    optimizer.zero_grad()
    logits = model(X)
    loss = criterion(logits, y)

    # ── Check 1: is the loss finite?
    if not torch.isfinite(loss):
        print(f"Step {step}: non-finite loss = {loss.item()}")
        print("Input contains NaN:", torch.isnan(X).any().item())
        print("Input contains Inf:", torch.isinf(X).any().item())
        break

    loss.backward()

    # ── Check 2: gradient norms — are gradients flowing at all?
    total_norm = sum(
        p.grad.norm().item() ** 2 for p in model.parameters() if p.grad is not None
    ) ** 0.5
    if step % 50 == 0:
        print(f"Step {step}: loss={loss.item():.4f} grad_norm={total_norm:.4f}")

    # ── Check 3: are any gradients None? (means that param was unused!)
    for name, p in model.named_parameters():
        if p.grad is None:
            print(f"WARNING: {name} has no gradient — is it used in forward()?")

    optimizer.step()

# ── Check 4: verify model output shape and range make sense
with torch.no_grad():
    sample_out = model(X[:1])
    print("Output range:", sample_out.min().item(), sample_out.max().item())

# ── Check 5: overfit a tiny batch — sanity check the architecture
# If the model cannot drive loss near zero on 5 examples, there is a bug
tiny_X, tiny_y = X[:5], y[:5]
for _ in range(200):
    optimizer.zero_grad()
    loss = criterion(model(tiny_X), tiny_y)
    loss.backward()
    optimizer.step()
print(f"Tiny-batch overfit loss: {loss.item():.6f}")  # should approach 0