How to add a tqdm progress bar to a PyTorch training loop

A training loop with an epoch bar and a batch bar

In a plain PyTorch training loop you want two bars: an outer one for epochs and an inner one wrapping your DataLoader for the batches. The current release is tqdm 4.70.1, and its README puts the overhead at about 60 nanoseconds per iteration. Wrap the epoch range with trange and the loader with tqdm, and pass leave=False to the inner bar so finished epochs don't pile up in your terminal.

from tqdm import tqdm, trange

for epoch in trange(num_epochs, desc="epochs"):
    model.train()
    for x, y in tqdm(train_loader, desc=f"epoch {epoch}", leave=False):
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        loss.backward()
        optimizer.step()

If you need to log something inside the loop, don't use a plain print(), because it breaks the bar. Use tqdm.write() instead, which prints a line above the bars without mangling them.

Show the loss with set_postfix

To put live metrics on the bar, call set_postfix() on the batch bar object. It adds key=value pairs after the rate, while set_description() changes the label on the left.

There's a hidden cost here. The PyTorch performance tuning guide lists cuda_tensor.item() among the operations that force a CPU-GPU synchronization. If you call loss.item() on every step, the CPU waits for the GPU to finish that batch before it can queue the next one. Accumulate the loss on the GPU and only pull it to the CPU every N steps.

LOG_EVERY = 50

bar = tqdm(train_loader, desc=f"epoch {epoch}", leave=False)
running = torch.zeros((), device=device)

for step, (x, y) in enumerate(bar, start=1):
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
    running += loss.detach()

    if step % LOG_EVERY == 0:
        bar.set_postfix(loss=f"{(running / LOG_EVERY).item():.4f}", lr=optimizer.param_groups[0]["lr"])
        running.zero_()

Nested bars: what position and leave do

By default tqdm uses leave=True and position=None, which means automatic placement. When you nest a batch bar inside an epoch bar, automatic placement puts the inner bar on the line below the outer one, and leave=True keeps a finished bar on screen when its loop ends.

That default is right for the outer epoch bar, since you want it visible when training finishes. On the inner bar, leave=True prints a new permanent line at the end of every epoch, which is why you set leave=False there so the next epoch reuses the same line. If automatic placement ever gets confused in your terminal, pass position=0 on the outer bar and position=1 on the inner one.

The tqdm settings that matter in training

Most of tqdm's arguments can stay at their defaults, but a handful matter for long training jobs or when you log to a file.

Multi-GPU and DDP: draw on rank 0 only

When you launch a DistributedDataParallel job, every process runs the training loop. Without disable, you get one bar per GPU fighting over the same terminal.

torchrun sets the environment variables RANK, LOCAL_RANK and WORLD_SIZE for every process. Read RANK and pass disable=rank != 0, so only rank 0 draws while the other processes train silently.

import os

rank = int(os.environ.get("RANK", 0))

for x, y in tqdm(train_loader, desc=f"epoch {epoch}", disable=rank != 0):
    ...

What the ETA actually means

The time block on a tqdm bar reads [elapsed<remaining, rate]. The remaining value is the number of items left divided by the current rate, and the rate is an exponential moving average with smoothing=0.3 by default.

Because of that moving average, the ETA jumps after a slow first batch, during validation or while a checkpoint is written. If the jumps bother you, lower smoothing toward 0 so the rate leans toward the average speed of the whole run.

Lightning and Keras do their own thing

If you use a framework instead of a plain loop, the bars work differently. In PyTorch Lightning 2.6.6, TQDMProgressBar(refresh_rate=1) is the tqdm bar, but RichProgressBar becomes the default when the rich package (10.2.2 or newer) is installed. Trainer(enable_progress_bar=False) turns bars off completely.

Keras doesn't use tqdm at all. It has its own Progbar, and fit(verbose=...) accepts "auto", 0, 1 or 2.

Watch the bar from your phone

The batch bar is the best signal that a run is healthy, so it's the thing you want to see when you're away from the terminal. Runpeek streams standard tqdm bars, including Lightning's TQDMProgressBar in a terminal; Rich bars and the Keras Progbar are not covered.

patch() makes every from tqdm import tqdm bar in the process stream, and start_run() groups them under one link. The rank check keeps a DDP job to a single link. Updates go out at most about every 0.75 seconds from a background thread, and network errors never raise into your training loop.