Как изменить модель Pytorch для получения данных от загрузчика данных

Question

Я делаю проект НЛП. Сначала я подаю предложения построчно в модель.

Я пытаюсь использовать DataLoader для одновременной подачи нескольких предложений в модель. Я реализовал этот код в DataSet:

    def __getitem__(self, i, cuda=use_cuda):
    # encode the sentence at index i
    sentence = self.corpus_original[i]
    sentence_lowered = self.corpus_lowered[i]
    sentence_tags = self.tags[i]

    max_sent_len = self.max_sent_len

    # Idea behind this is that we pad our arrays with zeros and create a new tag and word for those padded indices.
    words = torch.zeros(max_sent_len, dtype=torch.long)
    labels = torch.zeros(max_sent_len, dtype=torch.long)
    chars = torch.zeros(max_sent_len, self.max_word_len,
                        dtype=torch.long)
    word_lengths = np.zeros(max_sent_len, dtype=int)

    for j in range(len(sentence)):
        labels[j] = self.tag2idx[sentence_tags[j]]
        words[j] = self.word2idx[sentence_lowered[j]]
        word_lengths[j] = len(sentence[j])
        for k, c in enumerate(sentence[j]): # for each character
            chars[j,k] = self.char2idx[c]

    # pad here
    # we use -1 to denote padded values
    for j in range(len(sentence), max_sent_len):
        labels[j] = 45 # idx for <pad_tag>
        words[j] = 1
        word_lengths[j] = 5 # 5 because length of "<pad>" is 5
        chars[j, :] = 94 # 94 proper characters, 0 idx so 94 is our idx for <pad_char>

    if cuda:
        words = words.cuda()
        labels = labels.cuda()
        chars = chars.cuda()

    return words, labels, chars, word_lengths

Затем я создаю train_loader с помощью:

train_loader = DataLoader(train, batch_size=32, shuffle=False, num_workers=4)

Я генерирую каждую партию и пытаюсь позвонить

preds = model(words, chars, word_lengths)

Неудивительно, что он говорит мне, что

RuntimeError: Expected 3-dimensional input for 3-dimensional weight [10, 16, 5], but got input of size [16, 54, 141, 16] instead

У меня вопрос: как мне изменить код моей модели, чтобы она могла принимать пакет из 16 предложений вместо одного предложения?

Спасибо.