def new_arange(x, *size):
    """
    Return a Tensor of `size` filled with a range function on the device of x.
    If size is empty, using the size of the variable x.
    """
    if len(size) == 0:
        size = x.size()
    return torch.arange(size[-1], device=x.device).expand(*size).contiguous()

    Return a Tensor of `size` filled with a range function on the device of x.
    If size is empty, using the size of the variable x.
    
understand_icl/fairseq/fairseq/utils.py