Trainer accumulates GPU usage at the beginning of each step
Hello,

My problem is that GPU usage gets increased at the beginning of each step. Although the usage gets decreased with the help of torch.cuda.empty_cache() and gc.collector() during training, OOM errors happened after a while.

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.5.1
- Platform: colab
- Python version: 3.7.10
- PyTorch version (GPU?): 1.8.1+cu101
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No


### Who can help

@sgugger @patrickvonplaten 

<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Models I am using: wav2vec2 and MBart

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Build e2e_model with the following classes: wav2vec2_learn_repr and e2emodel
2. Feed audio and translation under the requirement of the following data_collator.
3. Model training with Trainer.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

The code for reproducing the error:
```
class wav2vec2_learn_repr(Wav2Vec2PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.wav2vec2 = Wav2Vec2Model(config)
        self.dropout = nn.Dropout(config.final_dropout)
        self.collapse = collapse_layer
        self.init_weights = ()

    def forward(self,
                input_values,
                attention_mask=None,
                output_attentions=None,
                output_hidden_states=None,
                return_dict=None,
                labels=None):

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.wav2vec2(
            input_values,
            attention_mask=attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = outputs[0]
        hidden_states = self.dropout(hidden_states)
        collapsed_embeddings, attention_masks=cal_collapse_embeddings(hidden_states)
        show_gpu(f'In wav2vec2_learn_repr before del')
        del outputs, hidden_states
        show_gpu(f'In wav2vec2_learn_repr after del')
        torch.cuda.empty_cache()
        show_gpu(f'In wav2vec2_learn_repr empty cache')
        gc.collect()
        show_gpu(f'In wav2vec2_learn_repr gc.collect()')
        return collapsed_embeddings, attention_masks
```
```
class e2emodel(PreTrainedModel):
  def __init__(self,
               wav2vec2_name = "facebook/wav2vec2-large-xlsr-53",  
               mbart_model_name = 'facebook/mbart-large-50-many-to-many-mmt', 
               ):
    super().__init__(PretrainedConfig())
    self.wav2vec2_repr_model = wav2vec2_learn_repr.from_pretrained(wav2vec2_name)
    self.mbart_model = MBartForConditionalGeneration.from_pretrained(mbart_model_name)
    self.wav2vec2_repr_model.to(device)
    self.mbart_model.to(device)
  
  def forward(self,
                input_ids,
                attention_mask=None,
                output_attentions=None,
                output_hidden_states=None,
                return_dict=None,
                labels=None):
    

    show_gpu(f'At start')
    torch.cuda.empty_cache()
    show_gpu(f'empty cache')
    gc.collect()
    show_gpu(f'gc.collect()')
    
    input_ids.to(device)
    # print(f' inputs devices: {input_ids.device}, {labels.device}')
    show_gpu(f'load input_ids')


    collapsed_embeddings, attention_masks = self.wav2vec2_repr_model(input_ids)

    print('collapsed_embeddings, attention_masks', collapsed_embeddings.device, attention_masks.device)
    show_gpu(f'after wav2vec2')
    del input_ids
    show_gpu(f'delete input_ids')
    torch.cuda.empty_cache()
    show_gpu(f'empty cache')
    gc.collect()
    show_gpu(f'gc.collect()')
    labels.to(device)
    show_gpu(f'load mbart inputs')

    output = self.mbart_model(inputs_embeds = collapsed_embeddings, attention_mask = attention_masks, labels = labels)

    show_gpu(f'after mbart')
    del collapsed_embeddings, attention_masks, labels
    show_gpu(f'delete mbart inputs')
    torch.cuda.empty_cache()
    show_gpu(f'empty cache')
    return output
```
```
def data_collator(data):
  translation = [d['translation'] for d in data]
  input_features = [{'input_values': get_inputs_values_from_audio_path(feature_extractor, d['path'])} for d in data]
  #TODO: remove empty audio and its translation
  wav2vec2_inputs = feature_extractor.pad(input_features,
                                padding=True,
                                max_length=None,
                                pad_to_multiple_of=None,
                                return_tensors="pt", )
  batch={}
  batch['inputs_embeds'], batch['attention_mask'] = wav2vec2_learn_repr(wav2vec2_inputs['input_values']) # size [batch_size, nr_sample, 1024]
  with tokenizer.as_target_tokenizer():
    batch['labels'] = tokenizer([d['translation']for d in data], return_tensors='pt', padding=True).input_ids
  return batch
```
```
import torchaudio
resampler = torchaudio.transforms.Resample(orig_freq=48000, new_freq=16000)
def get_inputs_values_from_audio_path(processor, path: str):
    signal, sr = torchaudio.load(main_path + '{}/clips/'.format(src_lang) + path)
    signal = signal.squeeze(0)
    d = (signal.shape[0]/sr)
    resampler.orig_freq = sr
    signal=resampler.forward(signal).numpy()
    input_values = processor(signal, sampling_rate=resampler.new_freq).input_values
    return input_values.tolist()[0]
```
```
import gc
import subprocess
def show_gpu(msg):
    """
    ref: https://discuss.pytorch.org/t/access-gpu-memory-usage-in-pytorch/3192/4
    """
    def query(field):
        return(subprocess.check_output(
            ['nvidia-smi', f'--query-gpu={field}',
                '--format=csv,nounits,noheader'], 
            encoding='utf-8'))
    def to_int(result):
        return int(result.strip().split('\n')[0])
    
    used = to_int(query('memory.used'))
    total = to_int(query('memory.total'))
    pct = used/total
    print('\n' + msg, f'{100*pct:2.1f}% ({used} out of {total})') 
```
## Expected behavior
<!-- A clear and concise description of what you would expect to happen. -->

Here is the GPU usage history from step 7. 

**At start 86.9% (14149 out of 16280)**
empty cache 54.3% (8835 out of 16280)
gc.collect() 54.3% (8835 out of 16280)
load input_ids 54.3% (8835 out of 16280)
In wav2vec2_learn_repr before del 56.7% (9231 out of 16280)
In wav2vec2_learn_repr after del 56.7% (9231 out of 16280)
In wav2vec2_learn_repr empty cache 56.7% (9229 out of 16280)
In wav2vec2_learn_repr gc.collect() 56.7% (9229 out of 16280)
collapsed_embeddings, attention_masks cuda:0 cuda:0
after wav2vec2 56.7% (9229 out of 16280)
delete input_ids 56.7% (9229 out of 16280)
empty cache 56.7% (9229 out of 16280)
gc.collect() 56.7% (9229 out of 16280)
load mbart inputs 56.7% (9229 out of 16280)
after mbart 64.3% (10473 out of 16280)
delete mbart inputs 64.3% (10473 out of 16280)
empty cache 64.3% (10473 out of 16280)
**At start 85.5% (13925 out of 16280)**
empty cache 54.3% (8835 out of 16280)
gc.collect() 54.3% (8835 out of 16280)
load input_ids 54.3% (8835 out of 16280)
In wav2vec2_learn_repr before del 58.9% (9593 out of 16280)
In wav2vec2_learn_repr after del 58.9% (9593 out of 16280)
In wav2vec2_learn_repr empty cache 58.9% (9593 out of 16280)
In wav2vec2_learn_repr gc.collect() 58.9% (9593 out of 16280)
collapsed_embeddings, attention_masks cuda:0 cuda:0
after wav2vec2 58.9% (9593 out of 16280)
delete input_ids 58.9% (9593 out of 16280)
empty cache 58.9% (9593 out of 16280
gc.collect() 58.9% (9593 out of 16280)
load mbart inputs 58.9% (9593 out of 16280)
after mbart 66.7% (10853 out of 16280)
delete mbart inputs 66.7% (10853 out of 16280)
empty cache 66.7% (10853 out of 16280)
**At start 95.4% (15529 out of 16280)**
empty cache 94.6% (15393 out of 16280)
gc.collect() 94.6% (15393 out of 16280)
load input_ids 94.6% (15393 out of 16280)
RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 15.90 GiB total capacity; 14.96 GiB already allocated; 21.75 MiB free; 15.00 GiB reserved in total by PyTorch)


I've experienced the same issue. 