image classification ```KeyError: 'image'```
Author: yongjer
Pull request: no

### System Info

- `transformers` version: 4.41.2
- Platform: Linux-6.5.0-35-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.23.3
- Safetensors version: 0.4.3
- Accelerate version: 0.30.1.dev0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.0a0+07cecf4168.nv24.05 (True)
- 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?: RTX 4060TI 16G
- Using distributed or parallel set-up in script?: no

### Who can help?

_No response_

### Information

- [X] 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

```
python ./run_image_classification.py \
--model_name_or_path google/vit-base-patch16-224 \
--train_dir dataset \
--do_train \
--do_predict \
--do_eval \
--eval_strategy steps \
--fp16 \
--torch_compile \
--torch_compile_mode max-autotune \
--output_dir /workspaces/image_cls/output \
--overwrite_output_dir
```
[run_image_classification.py](https://github.com/huggingface/transformers/blob/v4.41.2/examples/pytorch/image-classification/run_image_classification.py), but change ```#ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,``` to ```ignore_mismatched_sizes=True``` due to the error:
```
Traceback (most recent call last):
  File "/workspaces/image_cls/./run_image_classification.py", line 436, in <module>
    main()
  File "/workspaces/image_cls/./run_image_classification.py", line 312, in main
    model = AutoModelForImageClassification.from_pretrained(
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/auto/auto_factory.py", line 563, in from_pretrained
    return model_class.from_pretrained(
  File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py", line 3754, in from_pretrained
    ) = cls._load_pretrained_model(
  File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py", line 4265, in _load_pretrained_model
    raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")
RuntimeError: Error(s) in loading state_dict for ViTForImageClassification:
        size mismatch for classifier.weight: copying a param with shape torch.Size([1000, 768]) from checkpoint, the shape in current model is torch.Size([4, 768]).
        size mismatch for classifier.bias: copying a param with shape torch.Size([1000]) from checkpoint, the shape in current model is torch.Size([4]).
        You may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method.
```
#### here is my [dataset](https://www.kaggle.com/datasets/mahyeks/almond-varieties?select=dataset)

### Expected behavior

train successfully, here is the error message:
```
Traceback (most recent call last):
  File "/workspaces/image_cls/./run_image_classification.py", line 436, in <module>
    main()
  File "/workspaces/image_cls/./run_image_classification.py", line 410, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 1885, in train
    return inner_training_loop(
  File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 2178, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/usr/local/lib/python3.10/dist-packages/accelerate/data_loader.py", line 454, in __iter__
    current_batch = next(dataloader_iter)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py", line 631, in __next__
    data = self._next_data()
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py", line 675, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = self.dataset.__getitems__(possibly_batched_index)
  File "/usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py", line 2865, in __getitems__
    batch = self.__getitem__(keys)
  File "/usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py", line 2861, in __getitem__
    return self._getitem(key)
  File "/usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py", line 2846, in _getitem
    formatted_output = format_table(
  File "/usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py", line 633, in format_table
    return formatter(pa_table, query_type=query_type)
  File "/usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py", line 401, in __call__
    return self.format_batch(pa_table)
  File "/usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py", line 516, in format_batch
    return self.transform(batch)
  File "/workspaces/image_cls/./run_image_classification.py", line 361, in train_transforms
    _train_transforms(pil_img.convert("RGB")) for pil_img in example_batch[data_args.image_column_name]
KeyError: 'image'
  0%|                                                                                                                                     | 0/498 [00:00<?, ?it/s]
  ```

Comments:
Hi @yongjer, thanks for opening an issue! 

Two things: 
* It's not necessary to directly modify the code to allow for initializing with mismatched sizes. This can be controlled by passing `--ignore_mismatched_size True` when calling from the CLI
* The error is happening because the batch returned from the dataset doesn't have an `image` column. `--image_column_name` should be updated to the correct value for your dataset 

> Hi @yongjer, thanks for opening an issue!
> 
> Two things:
> 
> * It's not necessary to directly modify the code to allow for initializing with mismatched sizes. This can be controlled by passing `--ignore_mismatched_size True` when calling from the CLI
> * The error is happening because the batch returned from the dataset doesn't have an `image` column. `--image_column_name` should be updated to the correct value for your dataset

So what should I do? 
Now I have a folder called dataset, which has four subfolders, and each of them contains hundreds of images.
![螢幕快照 2024-06-05 23-56-18](https://github.com/huggingface/transformers/assets/54315206/e5afd953-3fcf-40bb-ae4e-579bc8e1a78a)


passing ```--remove_unused_columns False``` just works