Hi @BramVanroy,
Thanks for your quick reply. As I said, I am trying to implement a binary classification task, but the data is imbalanced. So  the weights that I used were self.hparams.class_weights = [1, 7.48] (in this list form).
The error is the following:
ValueError: Expected target size (8, 32128), got torch.Size([8, 2])

Since the logits tensor has shape (8, 2, 32128), and the labels tensor, (8,2). However this code worked perfectly like this:
def forward(
      self, input_ids, attention_mask=None, decoder_input_ids=None, decoder_attention_mask=None, lm_labels=None
  ):  # in lightning, forward defines the prediction/inference actions
    return self.model(
        input_ids,
        attention_mask=attention_mask,
        decoder_input_ids=decoder_input_ids,
        decoder_attention_mask=decoder_attention_mask,
        lm_labels=lm_labels
    )

  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']
    )
    loss = outputs[0]
    return loss

The problem is that in that way, the model learns to always return the majority class. That’s why I tried to implement another CrossEntropyLoss using the weights:
##### IMBALANCE LEARNING
    class_weights = torch.FloatTensor(self.hparams.class_weights).cuda()
    loss_fct = CrossEntropyLoss(ignore_index=-100, weight=class_weights)
    loss = loss_fct(logits, lm_labels)
    return loss

I hope I made it clearer now.
Marcos