Function reduce_tensor

cdslib/core/utils/multigpu_utils.py:150–154

Gather and sum the tensor from all gpus.

(tensor: torch.Tensor)

def reduce_tensor(tensor: torch.Tensor):
    """Gather and sum the tensor from all gpus."""
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    return rt