I’m using the ElectraModel.from_pretrained(‘google/electra-base-discriminator’) to train a multi label classification task. I would like to add a linear layer for the final hidden state logits. Below is a pseudocode example
model = ElectraModel.from_pretrained('google/electra-base-discriminator')
logit_layer = torch.nn.Linear(768, 4)
### below code is what I'm trying to figure out 
append_logit_layer = model.append(logit_layer)

My main reason for needing to append this is so when I back propagate with torch SGD
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.8)
The model.parameters() value will have my logit_layer gradient.