Mask when batching sequence


I want to get masked tensor when I batched the variable length sequences.

import torch.nn as nn

l = [torch.zeros(7, 4), torch.zeros(8, 4), torch.zeros(9, 4)]
b = nn.utils.rnn.pad_sequence(l, batch_first=True)
mask = # will be torch.BooleanTensor[3, 9] with true for valid input and false for padded input

Is there any simple implementation for this one?

mask = b[:, 0] != 0

This kind of things usually work in some case, but not all the case…