cf8a3fb30547d6e980ecd8182f64a51df8e55c62,python/dgl/backend/pytorch/tensor.py,,pack_padded_tensor,#Any#Any#,248

Before Change

def pack_padded_tensor(input, lengths):
    batch_size, max_len = input.shape[:2]
    device = input.device
    index = []
    for i, l in enumerate(lengths):
        index.extend(range(i * max_len, i * max_len + l))
    index = th.tensor(index).to(device)
    return gather_row(input.view(batch_size * max_len, -1), index)

def boolean_mask(input, mask):
    if "bool" not in str(mask.dtype):
        mask = th.tensor(mask, dtype=th.bool)

After Change

        lengths = th.tensor(lengths, dtype=th.int64, device=device)
    else:
        lengths = lengths.to(device)
    input = input.view(-1, *input.shape[2:])
    out_len = lengths.sum().item()
    index = th.ones(out_len, dtype=th.int64, device=device)
    cum_lengths = th.cumsum(lengths, 0)
    index[cum_lengths[:-1]] += (max_len - lengths[:-1])
    index = th.cumsum(index, 0) - 1
    return input[index]

def boolean_mask(input, mask):
    if "bool" not in str(mask.dtype):
        mask = th.tensor(mask, dtype=th.bool)

Instances