hey @bennicholl, one idea would be to subclass ElectraModel, add your logit_layer to the init and then override the forward method so the logit_layer is fed the outputs from the encoder (see source code 3).
there might be more elegant approaches, but i’m pretty sure this one should work