Hi @lewtun, I have something like this in mind.
def compute_loss(self,model,inputs,classifier):
    output = model.generate(inputs)
    text = tok.decode(output)
    # convert text to ids
    classifier_output = classifier(text)
    loss = loss_function(classifier_output, targets)
    return loss

Now, when I do loss.backward(), I want to update the weights of model and not classifier. I understand that there is non-differentiability involved as soon as I use the model to generate text, I want to know how to handle that?