Consider the following network snippet:
def __init__(self, model, n_class, dropout_rate,device):
        super(NewModel, self).__init__()
        
        self.bert = model
        self.linear = nn.Linear(self.bert.config.hidden_size, 2)
        self.linear_1 = nn.Linear(self.bert.config.hidden_size, self.bert.config.hidden_size)
        
        self.dropout_rate = dropout_rate
        self.dropout_1 = nn.Dropout(p = self.dropout_rate)

        self.activation = nn.LeakyReLU()
        
        self.bn = nn.BatchNorm1d(num_features = self.bert.config.hidden_size)
       
    def forward(self, batch):
        outputs = self.bert(
            input_ids = batch[0].to(self.device),
            attention_mask = batch[1].to(self.device),
            token_type_ids = None,
            position_ids = None,
            head_mask = None,
            inputs_embeds = None,
        )

        output = outputs[0]
        pooled_output=output[:,0]
        pooled_output = pooled_output.unsqueeze(0)

        pooled_output_1 = self.dropout_1(self.bn(pooled_output))

        logits = self.linear(F.leaky_relu(self.linear_1(pooled_output_1)))

My batch size is 16 and during training I get this error:
ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 16])

I already set drop_last=True in the DataLoader, but the error persists.
Any help would be greatly appreciated.