How to train a model on top of a transformer to output a sequence?
I think you should treat this problem as a binary classification problem.  For each word in the changed sentence, you will have a binary label: correct or incorrect.  I would recommend relabeling so that &quot;correct&quot; words will have a label of 0 and &quot;incorrect&quot; words will have a label of 1.  In your example you would have:
correct_sentence = &quot;we used to play together&quot;
changed_sentence = &quot;we use play to together&quot;
labels = [0, 1, 1, 1, 0]

And instead of padding with some special value, pad with the &quot;correct&quot; label (which would be 0 if you use my suggestion above).
Conventionally, class labels always start at index 0, so this labeling scheme will match what PyTorch expects for binary classification problems.

Next, you will need to change the activation function for your final Linear layer.  Right now, your model ends with just a Linear layer, meaning the output is unbounded.  This doesn't really make sense for classification problems, because you know that the output should always be in the range [0, C-1], where C is the number of classes.
Instead, you should apply an activation function to make your outputs behave more like class labels.  For a binary classification problem, a good choice for the final activation is torch.nn.Sigmoid.  You would modify your model definition like this:
class Camembert(torch.nn.Module):
    &quot;&quot;&quot;
    The definition of the custom model, last 15 layers of Camembert will be retrained
    and then a fcn to 512 (the size of every label).
    &quot;&quot;&quot;
    def __init__(self, cam_model):
        super(Camembert, self).__init__()
        self.l1 = cam_model
        total_layers = 199
        for i, param in enumerate(cam_model.parameters()):
            if total_layers - i &gt; hparams[&quot;retrain_layers&quot;]:
                param.requires_grad = False
            else:
                pass
        self.l2 = torch.nn.Dropout(hparams[&quot;dropout_rate&quot;])
        self.l3 = torch.nn.Linear(768, 512)
        self.activation = torch.nn.Sigmoid()

    def forward(self, ids, mask):
        _, output = self.l1(ids, attention_mask=mask)
        output = self.l2(output)
        output = self.l3(output)
        output = self.activation(output)
        return output

Your output will now have dimension (batch_size, 512, 1).  Each of the 512 outputs will be a number between 0 and 1.  You can treat this as the probability of each particular token being &quot;incorrect&quot;.  If the output is greater than 0.5, then the label becomes &quot;incorrect&quot;.  Otherwise, the label is &quot;correct&quot;.
Finally, since you're treating the problem as a binary classification problem, you will want to use binary cross-entropy loss (torch.nn.BCELoss).  Note that you will have to unsqueeze the labels to make their dimension match the dimension of your output.
model = Camembert(cam_model)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

input = &lt;tokenized, padded input sequence&gt;
labels = torch.tensor([0, 1, 1, 1, 0, . . .  , 0])
output = model(input)
loss = criterion(output, labels.unsqueeze(1))

optimizer.zero_grad()
loss.backward()
optimizer.step()