cusomcrossencodertrainer.py

# Wrapper class to re-define the loss function everywhen a model is
# re-generated for hyper-parameter search.
class CustomCrossEncoderTrainer(CrossEncoderTrainer):
    def __init__(self, *, loss_init=None, **kwargs):
        self.__loss_init = loss_init
        super().__init__(**kwargs)

    def call_model_init(self, trial=None):
        model = super().call_model_init(trial=trial)
        self.loss = self.__loss_init(model)
        return model