def sync_tensor(
    tensor: _torch.Tensor,
    reduce_op: _torch.distributed.ReduceOp = _torch.distributed.ReduceOp.AVG,
):
    """
    If the input tensor is not on CPU and torch.distributed is initialized,
    then sync the input tensor using the specified reduction operation across
    multiple GPUs and return the resulting tensor.
    Otherwise, return the original tensor.
    """

    if tensor.device != _torch.device("cpu") and _torch.distributed.is_initialized():
        world_size = _torch.distributed.get_world_size()
        if world_size >= 1:
            # Synchronous call
            _torch.distributed.all_reduce(tensor.data, op=reduce_op)

    return tensor

    If the input tensor is not on CPU and torch.distributed is initialized,
    then sync the input tensor using the specified reduction operation across
    multiple GPUs and return the resulting tensor.
    Otherwise, return the original tensor.
    
coremltools/optimize/torch/_utils/fsdp_utils.py