util.py 文件源码

def get_range_vector(size: int, is_cuda: bool) -> torch.Tensor:
    """
    Returns a range vector with the desired size, starting at 0. The CUDA implementation
    is meant to avoid copy data from CPU to GPU.
    """
    if is_cuda:
        indices = torch.cuda.LongTensor(size).fill_(1).cumsum(0) - 1
    else:
        indices = torch.arange(0, size).long()
    return Variable(indices, requires_grad=False)