ViltForTokenClassification not working for personalize multiclass classification. 
Author: matiasfreitas
Pull request: no

### System Info

I had several errors trying to use the ViLt code for multiclass classification.

On the lines 1473-1478 of the [modeling_vilt.py](https://github.com/huggingface/transformers/blob/v4.36.1/src/transformers/models/vilt/modeling_vilt.py#L1413) we have that code:

```
     if labels is not None:
            loss_fct = CrossEntropyLoss()
            # move labels to correct device to enable PP
            labels = labels.to(logits.device)
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
```

Based on my manual testing (I confess I'm not the most skilled to be sure about the theoretical correctness) and on the the file [modeling_vit.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/vit/modeling_vit.pyl)

I changed this lines to:
```
        loss = None
        if labels is not None:
            # move labels to correct device to enable model parallelism
            labels = labels.to(logits.device)
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                one_hot_labels = F.one_hot(labels, self.num_labels).float()
                loss = loss_fct(logits, one_hot_labels)
```
And the results are fine here in my computer.

I think that should be change on the library. 


### Who can help?

@ArthurZucker @amyeroberts @younesbelkada

### Information

- [X] 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

Running this code:

```
from transformers import DefaultDataCollator, TrainingArguments, Trainer
import evaluate
import numpy as np

from transformers import ViltForTokenClassification, TrainingArguments, Trainer
from torch import Tensor

def getModel(path = None):
    if(path is None):
        path = "dandelin/vilt-b32-finetuned-nlvr2"
        
    model = ViltForTokenClassification.from_pretrained(

        path,

        num_labels=len(label['label2idx']),

        id2label= label['idx2label'],

        label2id= label['label2idx'],

        return_dict=True,
        
        problem_type = "multi_label_classification"
    )

    return model

#Very simple data collator that simply collates batches of dict-like
#objects and performs special handling for potential keys named label and label_ids
data_collator = DefaultDataCollator()
accuracy = evaluate.load("accuracy")

def compute_metrics(eval_pred):

    predictions, labels = eval_pred

    predictions = np.argmax(predictions, axis=1)

    return accuracy.compute(predictions=predictions, references=labels)


training_args = TrainingArguments(
    
    output_dir=middle_step,                  
    # Directory where model checkpoints and logs will be saved.
    
    remove_unused_columns=False,             
    # Whether to remove unused columns from the input data before training.
    
    evaluation_strategy="epoch",             
    # The evaluation strategy to adopt during training. "epoch" evaluates at the end of each epoch.
    
    save_strategy="epoch",                   
    # The checkpoint save strategy during training. "epoch" saves at the end of each epoch.
    
    learning_rate=5e-5,                      
    # The initial learning rate for the optimizer.
    
    per_device_train_batch_size=16,          
    # Batch size per GPU or CPU for training.
    
    gradient_accumulation_steps=4,           
    # Gradient accumulation involves updating the model's weights 
    # only after accumulating gradients over multiple batches.
    # This can be useful when the effective batch size is too large to fit into GPU memory.
    # Instead of processing the entire batch at once, the model processes
    # smaller batches and accumulates gradients before updating the weights.
    
    per_device_eval_batch_size=16,           
    # Batch size per GPU or CPU for evaluation.
    
    num_train_epochs=6,                       
    # Total number of training epochs.
    
    warmup_ratio=0.1,                         
    # Ratio of total training steps used for warmup.
    
    logging_steps=10,                         
    # Log every n updates steps.
    
    load_best_model_at_end=True,             
    # Whether or not to load the best model found at the end of training.
    
    metric_for_best_model="accuracy",        
    # Metric used to determine the best model, e.g., "accuracy".
)


trainer = Trainer(

    model_init=getModel,

    args=training_args,

    data_collator=data_collator,

    train_dataset=dataset["train"],

    eval_dataset=dataset["val"],

    compute_metrics=compute_metrics,

)
```

### Expected behavior

Not raise a error on the line loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) because of the non match sizes of the tensor. 

Comments:
Hi @matiasfreitas 
Thanks for the issue! Can you past the full traceback of the issue here? 🙏 

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_133949/2760725053.py in <cell line: 84>()
     82 
     83 #The warning is going to be solved when fine tunning below
---> 84 trainer.train()
     85 trainer.save_model(final_step)

~/.local/lib/python3.10/site-packages/transformers/trainer.py in train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1553                 hf_hub_utils.enable_progress_bars()
   1554         else:
-> 1555             return inner_training_loop(
   1556                 args=args,
   1557                 resume_from_checkpoint=resume_from_checkpoint,

~/.local/lib/python3.10/site-packages/transformers/trainer.py in _inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   1858 
   1859                 with self.accelerator.accumulate(model):
-> 1860                     tr_loss_step = self.training_step(model, inputs)
   1861 
   1862                 if (

~/.local/lib/python3.10/site-packages/transformers/trainer.py in training_step(self, model, inputs)
   2723 
   2724         with self.compute_loss_context_manager():
-> 2725             loss = self.compute_loss(model, inputs)
   2726 
   2727         if self.args.n_gpu > 1:

~/.local/lib/python3.10/site-packages/transformers/trainer.py in compute_loss(self, model, inputs, return_outputs)
   2746         else:
   2747             labels = None
-> 2748         outputs = model(**inputs)
   2749         # Save past state if it exists
   2750         # TODO: this needs to be fixed and made cleaner later.

~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py 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 = [], []

~/.local/lib/python3.10/site-packages/transformers/models/vilt/modeling_vilt.py in forward(self, input_ids, attention_mask, token_type_ids, pixel_values, pixel_mask, head_mask, inputs_embeds, image_embeds, labels, output_attentions, output_hidden_states, return_dict)
   1476             # move labels to correct device to enable PP
   1477             labels = labels.to(logits.device)
-> 1478             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
   1479 
   1480         if not return_dict:

~/.local/lib/python3.10/site-packages/torch/nn/modules/module.py 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 = [], []

~/.local/lib/python3.10/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
   1172 
   1173     def forward(self, input: Tensor, target: Tensor) -> Tensor:
-> 1174         return F.cross_entropy(input, target, weight=self.weight,
   1175                                ignore_index=self.ignore_index, reduction=self.reduction,
   1176                                label_smoothing=self.label_smoothing)

~/.local/lib/python3.10/site-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3027     if size_average is not None or reduce is not None:
   3028         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3029     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
   3030 
   3031 

ValueError: Expected input batch_size (640) to match target batch_size (16).
```

I forgot to mention, but i changed the line 1472 to

`logits = self.classifier(sequence_output[:, 0,:])`

To be more similar to ViT model. 

And I believe ViltForImagesAndTextClassification can have the same problem, but I'm not sure about since I didn't tested this model nor knows well the goal of the model (reasoning, but only one per time?).

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.