Electra sequence classification with pytorch lightning issues with 'pooler_output'
I'm working on a sentence classification task with multiple binary labels attached to each sentence. I'm using Electra and pytorch lightning to do the job, but I've run into a problem. When I'm running the trainer.fit(model, data) I get the following error:
AttributeError: 'tuple' object has no attribute 'pooler_output'
The error is referring to line 13 in the section where I'm defining pl.LightningModule:
class CrowdCodedTagger(pl.LightningModule):

  def __init__(self, n_classes: int, n_training_steps=None, n_warmup_steps=None):
    super().__init__()
    self.electra = ElectraModel.from_pretrained(ELECTRA_MODEL_NAME, return_dict=False) #changed ElectraModel to ElectraForSequenceClassification
    self.classifier = nn.Linear(self.electra.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.electra(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output) # &lt;---- this is the line the error is referring to.
    output = torch.sigmoid(output)
    loss = 0
    if labels is not None:
        loss = self.criterion(output, labels)
    return loss, output

  def training_step(self, batch, batch_idx):
    input_ids = batch[&quot;input_ids&quot;]
    attention_mask = batch[&quot;attention_mask&quot;]
    labels = batch[&quot;labels&quot;]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log(&quot;train_loss&quot;, loss, prog_bar=True, logger=True)
    return {&quot;loss&quot;: loss, &quot;predictions&quot;: outputs, &quot;labels&quot;: labels}

  def validation_step(self, batch, batch_idx):
    input_ids = batch[&quot;input_ids&quot;]
    attention_mask = batch[&quot;attention_mask&quot;]
    labels = batch[&quot;labels&quot;]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log(&quot;val_loss&quot;, loss, prog_bar=True, logger=True)
    return loss

  def test_step(self, batch, batch_idx):
    input_ids = batch[&quot;input_ids&quot;]
    attention_mask = batch[&quot;attention_mask&quot;]
    labels = batch[&quot;labels&quot;]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log(&quot;test_loss&quot;, loss, prog_bar=True, logger=True)
    return loss

  def training_epoch_end(self, outputs):
    
    labels = []
    predictions = []
    for output in outputs:
      for out_labels in output[&quot;labels&quot;].detach().cpu():
        labels.append(out_labels)
      for out_predictions in output[&quot;predictions&quot;].detach().cpu():
        predictions.append(out_predictions)

    labels = torch.stack(labels).int()
    predictions = torch.stack(predictions)

    for i, name in enumerate(LABEL_COLUMNS):
      class_roc_auc = auroc(predictions[:, i], labels[:, i])
      self.logger.experiment.add_scalar(f&quot;{name}_roc_auc/Train&quot;, class_roc_auc, self.current_epoch)

  def configure_optimizers(self):

    optimizer = AdamW(self.parameters(), lr=2e-5)

    scheduler = get_linear_schedule_with_warmup(
      optimizer,
      num_warmup_steps=self.n_warmup_steps,
      num_training_steps=self.n_training_steps
    )

    return dict(
      optimizer=optimizer,
      lr_scheduler=dict(
        scheduler=scheduler,
        interval='step'
      )
    )

Can anyone point me in a direction to fix the error?
EXAMPLE OF DATA STRUCTURE (in CSV):
sentence                      label_1          label_2          label_3
Lorem ipsum dolor sit amet    1                0                1
consectetur adipiscing elit   0                0                0
sed do eiusmod tempor         0                1                1
incididunt ut labore et       1                0                0
Lorem ipsum dolor sit amet    1                0                1


ELECTRA has no pooler layer like BERT (compare the return section for further information).
In case you only want to use the [CLS] token for your sequence classification, you can simply take the first element of the last_hidden_state (initialize electra without return_dict=False):
output = self.classifier(output.last_hidden_state[:, 0])

