error with Vision Transformer (ViT)
Author: sunhaozhepy
Pull request: no

There seems to be a problem with the HuggingFace Vision Transformer: it takes up all the memory in GPU and renders it impossible to train the model.
When I was dealing with the inference just like in `https://huggingface.co/docs/transformers/model_doc/vit` with a single image, it works well. But when I tried to do some fine-tuning with the MNIST dataset intergrated in PyTorch, thus with batches, suddenly it doesn't work anymore. Some details concerning my problem:
1. PyTorch MNIST (torchvision.datasets.MNIST) is not stocked in common image file formats (e.g. jpg., jpeg., png.) which forbids me from using ImageFolder;
2. and as the common dataloader can't handle batches of images, I had to enter my custom transform (which uses feature_extractor) as a parameter when loading the dataset. Sot it appears to me that the feature extractor is handling one image at a time.
Here's my code:
```
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets
from transformers import ViTFeatureExtractor, ViTForImageClassification

feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')

class ImageTransform:
    def __init__(self):
        pass

    def __call__(self, image):
        output = torch.from_numpy(feature_extractor(images=image.convert("RGB")).pixel_values[0])
        return output

mnist_train_dataset = datasets.MNIST("/mnist_root", train=True, download=True, transform=ImageTransform())
mnist_test_dataset = datasets.MNIST("/mnist_root", train=False, download=True, transform=ImageTransform())
mnist_train_dataloader = DataLoader(mnist_train_dataset, batch_size=64, shuffle=True)
mnist_test_dataloader = DataLoader(mnist_test_dataset, batch_size=64, shuffle=False)

device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'
print(f"using {device}.")

model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k').to(device)

optimizer = optim.AdamW(model.parameters(), lr=1e-4)
loss_function = nn.CrossEntropyLoss()
num_epochs = 3

for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0
    for images, labels in mnist_train_dataloader:
        optimizer.zero_grad()
        images, labels = images.to(device), labels.to(device)
        outputs = model(pixel_values=images)
        loss = loss_function(outputs.logits, labels)
        epoch_loss += loss * len(images)
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch + 1}: Cross Entropy loss = {epoch_loss / len(mnist_train_dataset)}")
```
The error message:
```
Traceback (most recent call last):
  File "main.py", line 56, in <module>
    outputs = model(pixel_values=images)
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\transformers-4.6.1-py3.8.egg\transformers\models\vit\modeling_vit.py", line 603, in forward
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\transformers-4.6.1-py3.8.egg\transformers\models\vit\modeling_vit.py", line 507, in forward
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\transformers-4.6.1-py3.8.egg\transformers\models\vit\modeling_vit.py", line 346, in forward
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\transformers-4.6.1-py3.8.egg\transformers\models\vit\modeling_vit.py", line 278, in forward
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\transformers-4.6.1-py3.8.egg\transformers\models\vit\modeling_vit.py", line 221, in forward
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\transformers-4.6.1-py3.8.egg\transformers\models\vit\modeling_vit.py", line 165, in forward
RuntimeError: CUDA out of memory. Tried to allocate 38.00 MiB (GPU 0; 2.00 GiB total capacity; 1.63 GiB already allocated; 0 bytes free; 1.70 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
```
I thought at first that it was due to a too big batch_size. I turned it into `batch_size=2` which is of course to highlight the effect, and the error message becomes:
```
C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\cuda\Loss.cu:247: block: [0,0,0], thread: [0,0,0] Assertion `t >= 0 && t < n_classes` failed.
C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\cuda\Loss.cu:247: block: [0,0,0], thread: [1,0,0] Assertion `t >= 0 && t < n_classes` failed.
Traceback (most recent call last):
  File "main.py", line 59, in <module>
    loss.backward()
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\torch\_tensor.py", line 307, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "C:\Users\Sun\anaconda3\envs\squad\lib\site-packages\torch\autograd\__init__.py", line 156, in backward
    allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
RuntimeError: Unable to find a valid cuDNN algorithm to run convolution
```
And when I searched on Google, people is telling me that it was also due to a limit of allocation in GPU...

I really do hope that this can be solved as I can't quite figure out how to properly use the official model release of Google... Thank you very much.
BTW my GPU configuration:
```
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 512.15       Driver Version: 512.15       CUDA Version: 11.6     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ... WDDM  | 00000000:01:00.0 Off |                  N/A |
| N/A   36C    P0    N/A /  N/A |      0MiB /  2048MiB |      3%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+
```

Comments:
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.