Function lengths_to_mask

datasets/dataloader.py:183–190

(lengths: list[int],
                    device: torch.device,
                    max_len: int = None)

def lengths_to_mask(lengths: list[int],
                    device: torch.device,
                    max_len: int = None) -> torch.Tensor:
    lengths = torch.tensor(lengths, device=device)
    max_len = max_len if max_len else max(lengths)
    mask = torch.arange(max_len, device=device).expand(
        len(lengths), max_len) < lengths.unsqueeze(1)
    return mask