BERT без оболочки + слой GRU для анализа настроений

Question

Я пытаюсь настроить BERT без учета анализа настроений, добавив в конце один слой GRU. Проблема заключается в том, что выходная форма модели BERT не соответствует необходимой входной форме для слоя GRU.

У меня есть 2 класса: положительный и отрицательный, и размер пакета 25. BERT выводит тензор формы (25, 2). ГРУ нужен 3D ввод.

Код:

      class BERT_Sent(nn.Module):

    def __init__(self, bert, batch_size, input_size, hidden_size, num_layers, dropout=0.2):

      super(BERT_Sent, self).__init__()

      self.bert = bert
      weight = next(self.parameters()).data
      self.relu = nn.ReLU()
      self.softmax = nn.Softmax()
      self.h0 = weight.new(num_layers, batch_size, hidden_size).zero_()
      self.h0.requires_grad = True

      # GRU
      self.gru = nn.GRU(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, dropout=dropout) 

    def forward(self, token_ids, attention_masks, labels=None):
        out1 = self.bert(token_ids, 
                        token_type_ids = None, 
                        attention_mask = attention_masks,
                        labels = labels)

        out1 = self.relu(out1.logits)    #relu: felt cute might delete later
        out2, h1 = self.gru(out1, self.h0)
        return self.softmax(out2)

      optimizer = torch.optim.AdamW(model_.parameters(), lr = 2e-5, eps = 1e-08)

Также я не знаю, полезен ли h0, когда у меня есть только один слой GRU, поэтому я хотел бы, чтобы кто-нибудь объяснил. Я добавил «требуется град» на тот случай, если это можно изменить с помощью обучения.

Я использую эти параметры в качестве начальной пробы:

      model_ = BERT_Sent(model, batch_size, 15, 256, 1)

Пожалуйста, помогите мне понять, как правильно добавить слой GRU. Я никогда не работал ни с BERT, ни с RNN/GRU.