pytorch_padding_test.py

import torch
import torch.nn as nn
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
n_vocab = 3
n_layers = 1
n_embed = 5

def padding_test(x_padding):
    embedding = nn.Embedding(n_vocab, n_embed, padding_idx=x_padding)
    inputs = [[1,1,2,1,1,2,1],[1,1,2]]
    inputs = nn.utils.rnn.pad_sequence([torch.tensor(e, device=device) for e in inputs],
                                       batch_first=True,
                                       padding_value=x_padding)
    print(inputs)
    embedded = embedding(inputs)
    print(embedded)

# When using 0 for padding, PyTorch works well.
# However, this means that we cannot use 0 as the index of the word.
padding_test(0)

# When using -1 for padding, PyTorch raises "index out of range" error.
padding_test(-1)