Skip to content

← Devices and Data Movement step 3 of 4

Medium Primitives

Accumulate losses without extracting Python scalars

Jump to the task

Calling item() on a scalar tensor gives Python a number. For a CUDA tensor, that can make the host wait for device work needed to produce the value. Extracting a loss after every batch can reduce overlap between the host and device. The cost depends on the workload; it does not necessarily serialize all work on the whole device.

Keep a metric total as a tensor

For losses on the same device and with the same dtype, accumulate detached values and extract the result only when it is needed:

import torch

losses = [torch.tensor(1.0, requires_grad=True),
          torch.tensor(2.0, requires_grad=True)]
total = losses[0].new_zeros(())
for loss in losses:
    total = total + loss.detach()
print(total.requires_grad)  # False
print(total.item())         # 3.0; extraction happens once, after accumulation.

new_zeros makes the accumulator follow the first loss's device and dtype. Detaching is appropriate for a logging metric. If the total is the objective you plan to backpropagate, do not detach those losses: that would remove the gradient connections you need.

Tensor operations versus Python decisions

loss > threshold produces a boolean tensor and need not extract a Python value. Using that result in if loss > threshold requires a Python truth value and can wait for CUDA work. Similarly, item(), float(tensor) and int(tensor) extract a scalar. Formatting or printing tensor contents can also require values to reach the host.

Logging less often can reduce such waits. A progress bar does not inherently synchronize the device; it depends on whether preparing its displayed values requires tensor data from the device.

What this exercise measures

The harness counts the internal _local_scalar_dense dispatch used for scalar extraction. It runs on CPU, where those reads do not demonstrate a GPU wait. The counter is a narrow check for this extraction pattern, not proof that arbitrary code has no synchronization, data transfer or other overhead.

Your task

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

Sum a nonempty list of zero-dimensional floating-point tensors, all on the same device and with the same dtype. Return a detached scalar tensor on that device with that dtype. Do not extract Python numbers while accumulating.

For losses with values [1.0, 2.0, 3.0], return a tensor containing 6.0. The harness converts the final result to a number after it stops counting, rounds it to five decimals and expects zero scalar-extraction dispatches inside your function. The supplied cases use float32 CPU tensors.