_forward_unimplemented() got an unexpected keyword argument 'input_ids'
### System Info

 I am training on Google Colab Pro+ with following info:

- `transformers` version: 4.31.0
- Platform: Linux-5.15.109+-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- Accelerate version: 0.21.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.0.1+cu118 (False)
- Tensorflow version (GPU?): 2.12.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.7.0 (cpu)
- Jax version: 0.4.13
- JaxLib version: 0.4.13
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@sgugger
@ArthurZucker  
@younesbelkada

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Hi, I am trying to build a Multi-tasking Model following this [article](https://towardsdatascience.com/how-to-create-and-train-a-multi-task-transformer-model-18c54a146240) for two task: **zero-shot classification** **and sentiment analysis**.

**1. Creates the encoder and an output head for each task.**

I have try three different loaders based on the suggestion from the similar issues [here](https://github.com/huggingface/transformers/issues/21335): `AutoModel`, `AutoModelForSequenceClassification`, `RobertaForSequenceClassification `

```
import torch.nn as nn
from typing import List
from transformers import AutoModel, AutoModelForSequenceClassification, RobertaForSequenceClassification


class SequenceClassificationHead(nn.Module):
    def __init__(self, hidden_size, num_labels, dropout_p=0.1):
        super().__init__()
        self.num_labels = num_labels
        self.dropout = nn.Dropout(dropout_p)
        self.classifier = nn.Linear(hidden_size, num_labels)

        self._init_weights()

    def forward(self, sequence_output, pooled_output, labels=None, **kwargs):
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        loss = None
        if labels is not None:
            if labels.dim() != 1:
                # Remove padding
                labels = labels[:, 0]

            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(
                logits.view(-1, self.num_labels), labels.long().view(-1)
            )

        return logits, loss

    def _init_weights(self):
        self.classifier.weight.data.normal_(mean=0.0, std=0.02)
        if self.classifier.bias is not None:
            self.classifier.bias.data.zero_()


class MultiTaskModel(nn.Module):
    def __init__(self, encoder_name_or_path, tasks: List):
        super().__init__()

        self.encoder = AutoModel.from_pretrained(encoder_name_or_path)

        self.output_heads = nn.ModuleDict()
        for task in tasks:
            decoder = self._create_output_head(self.encoder.config.hidden_size, task)
            # ModuleDict requires keys to be strings
            self.output_heads[str(task.task_id)] = decoder

    @staticmethod
    def _create_output_head(encoder_hidden_size: int, task):
        if task.task_type == "seq_classification":
            return SequenceClassificationHead(encoder_hidden_size, task.num_labels)
        else:
            raise NotImplementedError()
```

**2. Define the metrics**

```
from transformers import EvalPrediction
import numpy as np
from datasets import load_metric
import evaluate

accuracy_metric = evaluate.load("accuracy")
f1_metric = evaluate.load("f1")
precision_metric = evaluate.load("precision")
recall_metric = evaluate.load("recall")

def compute_metrics(eval_preds: EvalPrediction):
    preds_dim = (eval_preds.predictions[0] if isinstance(eval_preds.predictions, tuple) else eval_preds.predictions).ndim

    if preds_dim == 2:
        # Sentiment analysis
        average="binary"
    elif preds_dim == 3:
        # Sequence classification
        average="macro"
    else:
        raise NotImplementedError()

    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)

    accuracy = accuracy_metric.compute(predictions=predictions, references=labels)["accuracy"]
    precision = precision_metric.compute(predictions=predictions, references=labels, average=average)["precision"]
    recall = recall_metric.compute(predictions=predictions, references=labels, average=average)["recall"]
    f1 = f1_metric.compute(predictions=predictions, references=labels, average=average)["f1"]

    return {"accuracy": accuracy, "precision": precision, "recall": recall, "f1": f1}
```

**3. Fine-tuning pre-trained model `vinai/phobert-base-v2` which based on `RoBERTa`**

```
from transformers import DataCollatorWithPadding, IntervalStrategy
from transformers.trainer_utils import get_last_checkpoint
import os
import random


transformers.logging.set_verbosity_info()


set_seed(42)

model_args = ModelArguments(encoder_name_or_path="vinai/phobert-base")

training_args = TrainingArguments(
      do_train=True,
      do_eval=True,
      output_dir="./mtl_zsl_sa_model",
      evaluation_strategy = IntervalStrategy.STEPS,
      per_device_train_batch_size=32,
      per_device_eval_batch_size=32,
      eval_steps = 500,
      save_steps = 2000,
      logging_steps = 500,
      learning_rate=5e-5,
      label_smoothing_factor=0.1,
      # fp16=True,
      num_train_epochs=3,
      weight_decay=0.01,
      save_strategy=IntervalStrategy.STEPS,
      load_best_model_at_end = True,
      metric_for_best_model = 'f1',
      optim="adamw_torch",
      resume_from_checkpoint=True,
      remove_unused_columns=False
  )

data_args = DataTrainingArguments(max_seq_length=128,
                                  max_train_samples=100)

tokenizer = AutoTokenizer.from_pretrained(
    model_args.encoder_name_or_path,
    cache_dir=model_args.cache_dir,
    use_fast=model_args.use_fast_tokenizer,
    revision=model_args.model_revision,
    use_auth_token=True if model_args.use_auth_token else None,
)

tasks, raw_datasets = load_datasets(tokenizer, data_args, training_args)

model = MultiTaskModel(model_args.encoder_name_or_path, tasks)

train_dataset = raw_datasets["train"]
eval_datasets = raw_datasets["validation"]

data_collator = DataCollatorWithPadding(
    tokenizer, pad_to_multiple_of=8
)

# Initialize our Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset
    compute_metrics=compute_metrics,
    tokenizer=tokenizer,
    data_collator=data_collator,
)

trainer.train()
```

This is what my merged tokenized datasets look like

```
DatasetDict({
    train: Dataset({
        features: ['input_ids', 'token_type_ids', 'attention_mask', 'task_ids'],
        num_rows: 29250
    })
    validation: [Dataset({
        features: ['input_ids', 'token_type_ids', 'attention_mask', 'task_ids'],
        num_rows: 1510
    }), Dataset({
        features: ['input_ids', 'token_type_ids', 'attention_mask', 'task_ids'],
        num_rows: 2262
    })]
})
```

This is my loaded model:

```
MultiTaskModel(
  (encoder): RobertaForSequenceClassification(
    (roberta): RobertaModel(
      (embeddings): RobertaEmbeddings(
        (word_embeddings): Embedding(64001, 768, padding_idx=1)
        (position_embeddings): Embedding(258, 768, padding_idx=1)
        (token_type_embeddings): Embedding(1, 768)
        (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (dropout): Dropout(p=0.1, inplace=False)
      )
      (encoder): RobertaEncoder(
        (layer): ModuleList(
          (0-11): 12 x RobertaLayer(
            (attention): RobertaAttention(
              (self): RobertaSelfAttention(
                (query): Linear(in_features=768, out_features=768, bias=True)
                (key): Linear(in_features=768, out_features=768, bias=True)
                (value): Linear(in_features=768, out_features=768, bias=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (output): RobertaSelfOutput(
                (dense): Linear(in_features=768, out_features=768, bias=True)
                (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
            )
            (intermediate): RobertaIntermediate(
              (dense): Linear(in_features=768, out_features=3072, bias=True)
              (intermediate_act_fn): GELUActivation()
            )
            (output): RobertaOutput(
              (dense): Linear(in_features=3072, out_features=768, bias=True)
              (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
              (dropout): Dropout(p=0.1, inplace=False)
            )
          )
        )
      )
    )
    (classifier): RobertaClassificationHead(
      (dense): Linear(in_features=768, out_features=768, bias=True)
      (dropout): Dropout(p=0.1, inplace=False)
      (out_proj): Linear(in_features=768, out_features=2, bias=True)
    )
  )
  (output_heads): ModuleDict(
    (0): SequenceClassificationHead(
      (dropout): Dropout(p=0.1, inplace=False)
      (classifier): Linear(in_features=768, out_features=2, bias=True)
    )
    (1): SequenceClassificationHead(
      (dropout): Dropout(p=0.1, inplace=False)
      (classifier): Linear(in_features=768, out_features=3, bias=True)
    )
  )
)
```

And I run into the following error:

```
[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1537             self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size
   1538         )
-> 1539         return inner_training_loop(
   1540             args=args,
   1541             resume_from_checkpoint=resume_from_checkpoint,

[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in _inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   1807 
   1808                 with self.accelerator.accumulate(model):
-> 1809                     tr_loss_step = self.training_step(model, inputs)
   1810 
   1811                 if (

[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in training_step(self, model, inputs)
   2652 
   2653         with self.compute_loss_context_manager():
-> 2654             loss = self.compute_loss(model, inputs)
   2655 
   2656         if self.args.n_gpu > 1:

[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in compute_loss(self, model, inputs, return_outputs)
   2677         else:
   2678             labels = None
-> 2679         outputs = model(**inputs)
   2680         # Save past state if it exists
   2681         # TODO: this needs to be fixed and made cleaner later.

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks
   1500                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501             return forward_call(*args, **kwargs)
   1502         # Do not call functions when jit is used
   1503         full_backward_hooks, non_full_backward_hooks = [], []

TypeError: _forward_unimplemented() got an unexpected keyword argument 'input_ids'
```

### Expected behavior

The trainer run smoothly without errors
I have add another `forward` method for my `MultiTaskModel `and the error is gone even though the new one appear. 