As far as how to obtain the mask, I’m not sure if this has been addressed, or if there is a better/easier way, but what I’ve been doing is something like:
cumsum = torch.Tensor(N, K).fill_(1).cumsum(-1) - 1
mask = cumsum < lens.view(N, 1).expand_as(cumsum)

Not sure if there is a better/easier/more standard way of creating the mask?