Two question about add a pretrained BERT in Text classification example as a layer

ttm: <p>Hi everyone<br>
I am trying to add a pretrained BERT in Text classification example as a layer.<br>
Firstly, the example set batch_size as 16, so I got the following error:</p>
<pre><code class="lang-auto">ValueError: Expected input batch_size(1) to match target batch_size(16)
</code></pre>
<p>Here is my code about model and train function</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch.nn.functional as F
import numpy
class TextSentiment(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_class):
        super().__init__()
        self.bert = BertModel.from_pretrained(model_name)
        self.fc = nn.Linear(embed_dim, num_class)

    def forward(self, text, offsets):        
        outputs = self.bert(text)
        last_hidden_states = outputs[0]
        print(last_hidden_states.shape) 
        
        return self.fc(last_hidden_states)
</code></pre>
<pre><code class="lang-auto">from sklearn.metrics import f1_score
from torch.utils.data import DataLoader

def train_func(sub_train_):

    # Train the model
    train_loss = 0
    train_acc = 0
    train_f1 = 0
    data = DataLoader(sub_train_, batch_size=BATCH_SIZE, shuffle=True,
                      collate_fn=generate_batch)
    for i, (text, offsets, cls) in enumerate(data):
        optimizer.zero_grad()
        offsets, cls = offsets.to(device), cls.to(device)
        input_text = r"C:\pyproject\.data\ag_news_csv\train.csv"
        text = berttokenizer.encode(input_text, add_special_tokens=True)
        text = torch.LongTensor([text])
        print(1, text.shape)
        print(text)
        
        output = model(text, offsets)
        print(2, output.shape)
        print(output)
        loss = criterion(output, cls)
        train_loss += loss.item()
        loss.backward()
        optimizer.step()
        train_f1 += f1_score(cls, output.argmax(1), average='macro')
        train_acc += (output.argmax(1) == cls).sum().item()

    # Adjust the learning rate
    scheduler.step()

    return train_loss / len(sub_train_), train_acc / len(sub_train_), train_f1 / len(data) 
</code></pre>
<p>and here is the result<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/c/eccb55b42162af75fdc8504469b30eb015a7d08a.png" data-download-href="https://discuss.pytorch.org/uploads/default/eccb55b42162af75fdc8504469b30eb015a7d08a" title="e08952127ec727405cb54c35859e8e0"><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/c/eccb55b42162af75fdc8504469b30eb015a7d08a.png" alt="e08952127ec727405cb54c35859e8e0" data-base62-sha1="xMMkcj6TDTlXi0HbcX9kKm0rNZo" width="483" height="499" data-dominant-color="2D2D2E"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">e08952127ec727405cb54c35859e8e0</span><span class="informations">805×833 29 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I do not know how to let input batch_size as 16, so I change the batch_size as 1, and I got another problem:</p>
<pre><code class="lang-auto">ValueError: Expected target size (1,4), got torch.Size([1])
</code></pre>
<p>Obviously, my problem is in the data shape but I can not fix it because I am new to pytorch or python.<br>
Any help will be appreciated.</p>

ttm: <p>I have already fixed the bug two</p>
<p>outputs[0] is the all embedding, and I only need cls vector, which means I need outputs[1].  And it worked!</p>
<p>But bug one still exist, and I am trying to rewrite a customer dataset and dataloader instead of example’s.  Hope it works.</p>