torch.nn.utils.rnn.pad_packed_sequence (original) (raw)

Pad a packed batch of variable length sequences.

The returned Tensor’s data will be of size T x B x * (if batch_first is False) or B x T x * (if batch_first is True) , where T is the length of the longest sequence and B is the batch size.

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence seq = torch.tensor([[1, 2, 0], [3, 0, 0], [4, 5, 6]]) lens = [2, 1, 3] packed = pack_padded_sequence( ... seq, lens, batch_first=True, enforce_sorted=False ... ) packed PackedSequence(data=tensor([4, 1, 3, 5, 2, 6]), batch_sizes=tensor([3, 2, 1]), sorted_indices=tensor([2, 0, 1]), unsorted_indices=tensor([1, 2, 0])) seq_unpacked, lens_unpacked = pad_packed_sequence(packed, batch_first=True) seq_unpacked tensor([[1, 2, 0], [3, 0, 0], [4, 5, 6]]) lens_unpacked tensor([2, 1, 3])