Quick tour runs into OOM on Colab 
Author: SaschaHeyer
Pull request: no

## Environment info
Google Colab

### Who can help
@jplu @LysandreJik @sgugger

## Information
Followed the quick tour using a Colab notebook https://huggingface.co/docs/datasets/quicktour.html#fine-tuning-a-deep-learning-model
Colab Runtime type: GPU

But the training process runs into OOM

```
---------------------------------------------------------------------------
ResourceExhaustedError                    Traceback (most recent call last)
<ipython-input-7-cb1582039f5e> in <module>()
      2 opt = tf.keras.optimizers.Adam(learning_rate=3e-5)
      3 model.compile(optimizer=opt, loss=loss_fn, metrics=["accuracy"])
----> 4 model.fit(tfdataset, epochs=3)

6 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     58     ctx.ensure_initialized()
     59     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 60                                         inputs, attrs, num_outputs)
     61   except core._NotOkStatusException as e:
     62     if name is not None:

ResourceExhaustedError: 2 root error(s) found.
  (0) Resource exhausted:  OOM when allocating tensor with shape[32,512,768] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
	 [[node tf_bert_for_sequence_classification/bert/encoder/layer_._4/attention/output/LayerNorm/batchnorm/mul_1 (defined at /usr/local/lib/python3.6/dist-packages/transformers/models/bert/modeling_tf_bert.py:327) ]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

	 [[gradient_tape/tf_bert_for_sequence_classification/bert/embeddings/position_embeddings/embedding_lookup/Reshape/_532]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

  (1) Resource exhausted:  OOM when allocating tensor with shape[32,512,768] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
	 [[node tf_bert_for_sequence_classification/bert/encoder/layer_._4/attention/output/LayerNorm/batchnorm/mul_1 (defined at /usr/local/lib/python3.6/dist-packages/transformers/models/bert/modeling_tf_bert.py:327) ]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_24759]

Function call stack:
train_function -> train_function
```

## To reproduce

Steps to reproduce the behavior:
The steps to reproduce are in the public accessible Colab notebook
[https://colab.research.google.com/drive/1Q3tBx57f2A8Hn1D7IXS-1nKandB86S3f](https://colab.research.google.com/drive/1Q3tBx57f2A8Hn1D7IXS-1nKandB86S3f)

## Expected behavior

Sample runs


Comments:
If you don't have any memory left, you should use a lower batch size. In the line:
``` 
>>> tfdataset = tf.data.Dataset.from_tensor_slices((features, dataset["labels"])).batch(32)
```
replace 32 by something lower.

Also, please use the [forums](https://discuss.huggingface.co/) for this kind of questions.