Are we able to create a custom class with extra layers and dropout when using trainer api?

Hi guys, I was looking at code that people used when the trainer api wasn’t and was included and it seems like people don’t create custom classes for their models anymore and just use the pre-trained one. For example I am talking about this class:

class SentimentClassifier(nn.Module):

  def __init__(self, n_classes):
    super(SentimentClassifier, self).__init__()
    self.bert = BertModel.from_pretrained(PRE_TRAINED_MODEL_NAME)
    self.drop = nn.Dropout(p=0.3)
    self.out = nn.Linear(self.bert.config.hidden_size, n_classes)
  
  def forward(self, input_ids, attention_mask):
    _, pooled_output = self.bert(
      input_ids=input_ids,
      attention_mask=attention_mask
    )
    output = self.drop(pooled_output)
    return self.out(output)

model = SentimentClassifier(2) 

source: Google Colab

And when I try to have my trainer use that model, I get this error from trainer.train()

RuntimeError: grad can be implicitly created only for scalar outputs

I tried looking around and I don’t see a solution for this. Is there someone out there who coded a tutorial on using classes that loads a pretrained model and add some layers on top of it and ran it inside the trainer api? Nowadays, it seems to me like we should just use model = BertForSequenceClassification.from_pretrained(model_name, num_labels=2) without extra layers on top which is fine by me. I just wanted to double check with you guys.

Thank you!