Hi @lewtun & @sgugger ! I have been having a similar issue with compute_loss () for  WeightedLossTrainer. Here is the code for the WeightedLossTrainer
class WeightedLossTrainer(Trainer):
      def compute_loss(self, model ,inputs, return_outputs=False):
           outputs = model(**inputs)
           logits = outputs.get("logits")
           labels= inputs.get("labels")
           loss_func=torch.nn.CrossEntropyLoss(weight=TorchWeights)
           loss = loss_func(logits,labels)  
           return (loss(logits, labels), outputs) if return_output else loss(logits, labels)

This is the code for tokenizing and aligning labels
def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(
        examples["tokens"], truncation=True,padding=True,is_split_into_words=True,return_offsets_mapping=True
)
    all_labels = examples["ner_tags"]
    new_labels = []
    for i, labels in enumerate(all_labels):
        word_ids = tokenized_inputs.word_ids(i)
        new_labels.append(align_labels_with_tokens(labels, word_ids))

    tokenized_inputs["labels"] = new_labels
    return tokenized_inputs

Here is a screenshot of the error :

Error_WL1721×339 66.6 KB

Any help?  I can share the full code if needed.
Best,
Ghadeer