AttributeError: 'SequenceClassifierOutput' object has no attribute 'pooler_output'

I am training the model Distilbert model for multi-label sequence classification.

from transformers import DistilBertTokenizer, DistilBertModel
import torch

tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertModel.from_pretrained('distilbert-base-uncased',return_dict=True)
n_classes=100
criterion = nn.BCELoss()
Batch_size=16

class SRTagger(pl.LightningModule):

  def __init__(self, n_classes: int, n_training_steps=None, n_warmup_steps=None):
    super().__init__()
    self.save_hyperparameters()
    self.bert = DistilBertModel.from_pretrained('distilbert-base-uncased',return_dict=True)
    self.classifier = nn.Linear(self.bert.config.hidden_size, n_classes)
    self.n_training_steps = n_training_steps
    self.n_warmup_steps = n_warmup_steps
    self.criterion = nn.BCELoss()

  def forward(self, input_ids, attention_mask, labels=None):
    output = self.bert(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output)
    output = torch.sigmoid(output)    
    loss = 0
    if labels is not None:
        loss = self.criterion(output, labels)
    return loss, output

It throws me the error:

AttributeError: ‘SequenceClassifierOutput’ object has no attribute ‘pooler_output’

Can I make changes to above code to overcome this?

[NielsRogge]

You can pass it when initializing the model (sorry should have mentioned that), like so:

from transformers import DistilBertForSequenceClassification

model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", problem_type="multi_label_classification",
num_labels=10)

[pratikchhapolika]

Thanks. Closing!!