Issue doing multi gpu training with TrOCR Transformer
Author: radu-diaconescu13
Pull request: no

Hello.

Thank you for your work on this lovely product.

I have managed to train the model on my dataset using the awesome notebook provided by @NielsRogge 
[here](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/TrOCR/Fine_tune_TrOCR_on_IAM_Handwriting_Database_using_native_PyTorch.ipynb)

I am having some issues doing multi gpu training on a single node. i.e using DataParallel

The way I set up my model is the following:


```
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-printed")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = DataParallel(model)

// do this because of previous issue with hugginface model with multi gpu
model = model.module

model.to(device)
```


and later on, in the for loop, i use the same device for the inputs:
```
  for i, (images, labels) in enumerate(data_loader):

      images, labels = images.to(device), labels.to(device)
```


Nevertheless, I am still getting:

```
  File "/opt/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/miniconda/lib/python3.7/site-packages/transformers/models/vision_encoder_decoder/modeling_vision_encoder_decoder.py", line 463, in forward
    **kwargs_encoder,
  File "/opt/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/miniconda/lib/python3.7/site-packages/transformers/models/vit/modeling_vit.py", line 557, in forward
    return_dict=return_dict,
  File "/opt/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/miniconda/lib/python3.7/site-packages/transformers/models/vit/modeling_vit.py", line 388, in forward
    layer_outputs = layer_module(hidden_states, layer_head_mask, output_attentions)
  File "/opt/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/miniconda/lib/python3.7/site-packages/transformers/models/vit/modeling_vit.py", line 319, in forward
    output_attentions=output_attentions,
  File "/opt/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/miniconda/lib/python3.7/site-packages/transformers/models/vit/modeling_vit.py", line 262, in forward
    self_outputs = self.attention(hidden_states, head_mask, output_attentions)
  File "/opt/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/miniconda/lib/python3.7/site-packages/transformers/models/vit/modeling_vit.py", line 191, in forward
    attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
RuntimeError: CUDA out of memory. Tried to allocate 244.00 MiB (GPU 0; 11.17 GiB total capacity; 10.54 GiB already allocated; 85.44 MiB free; 10.63 GiB reserved in total by PyTorch)

```


Any help is greatly appreciated.

Thank you!

Comments:
Hi,

Thanks for your interest in TrOCR! What's the batch size you are using?



Hi @NielsRogge 

Currently trying with batch size 16 on 4 Nvidia K80 gpu's on an Azure Compute.

It previously worked with batch size 4 on 1 Nvidia k80


Hello @NielsRogge 

Any idea about this?

I have managed to add automatic mixed precision training from pytorch on V100 compute
but I can only get up to batch size 8 this way.

Any help would be greatly appreciated.

Thank you

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/master/CONTRIBUTING.md) are likely to be ignored.