VideoMAE pretraining error when customizing compute_metrics
Author: gindij
Pull request: no

### System Info

- `transformers` version: 4.29.2
- Platform: macOS-13.4.1-arm64-i386-64bit
- Python version: 3.10.11
- Huggingface_hub version: 0.15.1
- Safetensors version: not installed
- PyTorch version (GPU?): 2.0.1 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help?
@sgugger @amyeroberts

### Information

- [ ] The official example scripts
- [ ] 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
I am trying to pre-train a VideoMAE model with a custom set of videos, and have found that by default, the evaluation loss is not reported after each epoch.

When I try to run this code 
```py
config = VideoMAEConfig.from_pretrained("MCG-NJU/videomae-base")
videomae = VideoMAEForPreTraining(config=config)

# <load datasets>

trainer = Trainer(
        model=videomae,
        args=training_arguments,
        optimizers=(torch.optim.AdamW(videomae.parameters()), None),
        data_collator=lambda x: _data_collator(x, videos_only=True, model_config=config),
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=_compute_metrics_videomae,
)

trainer.train()
```

I get this error message 
```
raceback (most recent call last):                                                                                                                                                                                       | 0/1 [00:00<?, ?it/s]
  File "/Users/jackgindi/Projects/echo-gpt/runtask.py", line 220, in <module>
    pretrain_videomae(
  File "/Users/jackgindi/Projects/echo-gpt/training.py", line 261, in pretrain_videomae
    trainer.train()
  File "/Users/jackgindi/miniconda3/envs/echo-gpt/lib/python3.10/site-packages/transformers/trainer.py", line 1664, in train
    return inner_training_loop(
  File "/Users/jackgindi/miniconda3/envs/echo-gpt/lib/python3.10/site-packages/transformers/trainer.py", line 2034, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/Users/jackgindi/miniconda3/envs/echo-gpt/lib/python3.10/site-packages/transformers/trainer.py", line 2300, in _maybe_log_save_evaluate
    metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
  File "/Users/jackgindi/miniconda3/envs/echo-gpt/lib/python3.10/site-packages/transformers/trainer.py", line 3029, in evaluate
    output = eval_loop(
  File "/Users/jackgindi/miniconda3/envs/echo-gpt/lib/python3.10/site-packages/transformers/trainer.py", line 3305, in evaluation_loop
    all_preds = nested_truncate(all_preds, num_samples)
  File "/Users/jackgindi/miniconda3/envs/echo-gpt/lib/python3.10/site-packages/transformers/trainer_pt_utils.py", line 357, in nested_truncate
    return type(tensors)(nested_truncate(t, limit) for t in tensors)
  File "/Users/jackgindi/miniconda3/envs/echo-gpt/lib/python3.10/site-packages/transformers/trainer_pt_utils.py", line 357, in <genexpr>
    return type(tensors)(nested_truncate(t, limit) for t in tensors)
  File "/Users/jackgindi/miniconda3/envs/echo-gpt/lib/python3.10/site-packages/transformers/trainer_pt_utils.py", line 361, in nested_truncate
    return tensors[:limit]
IndexError: too many indices for array: array is 0-dimensional, but 1 were indexed
```

If I run the same code without the `compute_metrics=...` set in the trainer, I don't encounter an error because `nested_truncate` is not called. The issue does not seem to with my my compute_metrics function, since the error occurs even before entering it. Is this just a current limitation of the VideoMAEForPreTraining model, or is there a way around this?

### Expected behavior

Pass a custom `compute_metrics` function to the `Trainer` to see the evaluation loss of `VideoMAEForPreTraining` after each epoch without error.

Comments:
Hi @gindij, 

For us help debug the issue, it's necessary to be able to reproduce the error on our side. At the moment this isn't possible without knowing `_data_collator`, `_compute_metrics_videomae` or the dataset. Could you share a minimal reproducer please? 

@amyeroberts thanks for the quick reply!

Here is a minimal reproducer:
```py
from typing import List

import datasets
import torch
from transformers import (
    EvalPrediction,
    PretrainedConfig,
    Trainer,
    TrainingArguments,
    VideoMAEConfig,
    VideoMAEForPreTraining,
)


def compute_image_mask(model_config: PretrainedConfig) -> torch.tensor:
    num_patches_per_frame = (model_config.image_size // model_config.patch_size) ** 2
    seq_length = (model_config.num_frames // model_config.tubelet_size) * num_patches_per_frame
    p = torch.ones((1, seq_length)) * (1 - model_config.masking_ratio)
    return torch.bernoulli(p).bool()


def data_collator(
    batch: List[dict],
    model_config: PretrainedConfig = None,
) -> dict:
    padded_videos = [torch.Tensor(item["video"]) for item in batch]
    padded_videos = torch.stack(padded_videos)
    mask = compute_image_mask(model_config)
    mask = mask.repeat((padded_videos.shape[0], 1))
    return {
        "pixel_values": padded_videos,
        "bool_masked_pos": mask,
    }


def compute_metrics_videomae(eval_pred: EvalPrediction) -> dict:
    pass


if __name__ == "__main__":
    config = VideoMAEConfig.from_pretrained("MCG-NJU/videomae-base")
    config.num_frames = 32
    config.masking_ratio = 0.9
    config.tubelet_size = 4
    config.patch_size = 32
    videomae = VideoMAEForPreTraining(config=config)

    train_dataset = {"video": [torch.rand((32, 3, 224, 224)) for _ in range(8)]}
    eval_dataset = {"video": [torch.rand((32, 3, 224, 224)) for _ in range(8)]}

    train_dataset = datasets.Dataset.from_dict(train_dataset)
    eval_dataset = datasets.Dataset.from_dict(eval_dataset)

    training_arguments = TrainingArguments(
        output_dir="./checkpts",
        per_device_eval_batch_size=8,
        per_device_train_batch_size=8,
        remove_unused_columns=False,
        evaluation_strategy="epoch",
        num_train_epochs=1,
    )

    trainer = Trainer(
        model=videomae,
        args=training_arguments,
        optimizers=(torch.optim.AdamW(videomae.parameters()), None),
        data_collator=lambda x: data_collator(x, config),
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics_videomae,  # <--- comment this line for it to work
    )

    trainer.train()
```

@amyeroberts, do you know of any updates on this issue?

@gindij In the example script, does `compute_metrics_videomae` intentionally return `None`? This is the cause of the error at the moment, as `Trainer` expects `compute_metrics` to be a callable which returns a dictionary. 

`eval_loss` isn't returned in the 'normal' loop because training an MAE pretraining model is a special case, as there are no labels passed in and it doesn't have a `return_loss` flag in its forward method. The simplest way to get what you want is to use your own custom training loop e.g. [similar to this one](https://github.com/huggingface/transformers/blob/main/examples/pytorch/image-pretraining/run_mim_no_trainer.py).

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.