The metric that stops the pipeline

total = 0.0
for batch in loader:
    loss = compute(batch)
    total += loss.item()          # once per batch

Correct, and it serialises the whole training loop.

Why

CUDA work is asynchronous. Python queues kernels and moves on; the device runs them at its own pace. That is what lets the host prepare batch n + 1 while the GPU is still on batch n.

.item() needs an actual number, so it has to wait for every queued kernel to finish. One .item() per batch means the host and device take turns instead of overlapping, and the pipeline is exactly as fast as the slower one doing all the work alone.

It is visible in the dispatch trace as aten._local_scalar_dense, which is the operation “read one element back to the host”. float(t), int(t) and an if on a tensor all lower to it.

The fix

Accumulate on the device and read once:

total = torch.zeros((), device=device)
for batch in loader:
    total += compute(batch).detach()
print(total.item())               # one sync, at the end

detach so the running total does not keep every batch’s graph alive, which is the autograd track’s lesson doing real work here.

The ones that surprise people

if loss > threshold: ...        # a sync
print(f"loss: {loss}")          # a sync
losses.append(loss.item())      # a sync
tqdm.set_postfix(loss=...)      # a sync, once a step

Progress bars showing a live loss are the most common accidental version. Logging every N steps rather than every step is usually enough, and costs nothing the rest of the time.

On CPU

There is nothing to wait for, so the cost is small. Write it correctly anyway: the same code moves to a GPU unchanged, and this is one of the few performance habits that costs nothing to adopt early.

Your task

def total_loss(losses: list) -> torch.Tensor

Sum a list of 0-dimensional loss tensors, performing no host synchronisations, and return the total as a tensor.