Hi @BramVanroy,
IMPORTANT UPDATE:
I have trying different things. First of all, I have checked that the loss produced by the model is the same as the CrossEntropy one without class_weights:
  def _step(self, batch):
    lm_labels = batch["target_ids"]
    lm_labels[lm_labels[:, :] == self.tokenizer.pad_token_id] = -100

    outputs = self(
        input_ids=batch["source_ids"],
        attention_mask=batch["source_mask"],
        lm_labels=lm_labels,
        decoder_attention_mask=batch['target_mask']
    )
    print("Loss 1", outputs[0])
    logits = outputs[1]
    loss_fct = CrossEntropyLoss(ignore_index=-100)
    loss = loss_fct( logits.view(-1, logits.size(-1)), lm_labels.view(-1))
    print("Loss 2", loss)
    return loss

This works, perfectly. However, when I introduce class_weights like this:
  def _step(self, batch):
    lm_labels = batch["target_ids"]
    lm_labels[lm_labels[:, :] == self.tokenizer.pad_token_id] = -100

    outputs = self(
        input_ids=batch["source_ids"],
        attention_mask=batch["source_mask"],
        lm_labels=lm_labels,
        decoder_attention_mask=batch['target_mask']
    )
    print("Loss 1", outputs[0])
    logits = outputs[1]
    class_weights = torch.FloatTensor([1, 7.8]).cuda()
    loss_fct = CrossEntropyLoss(ignore_index=-100, weight=class_weights)
    loss = loss_fct( logits.view(-1, logits.size(-1)), lm_labels.view(-1))
    print("Loss 2", loss)
    return loss

I get the following error:
RuntimeError: weight tensor should be defined either for all 32128 classes or no classes but got weight tensor of shape: [2]

Why is this happening? Since it is a binary classification problem.
Best,
Marcos