Как пакетно конвертировать длины предложений в маски в PyTorch?

Question

Например, из

lens = [3, 5, 4]

мы хотим получить

mask = [[1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 0]]

Оба из которых torch.LongTensor с.

Answer 1

torch.arange(max_len)[None, :] < lens[:, None]

Answer 2

Один из способов, который я нашел:

torch.arange(max_len).expand(len(lens), max_len) < lens.unsqueeze(1)

Пожалуйста, поделитесь, если есть лучшие способы!