Models fine-tuned with gradient checkpointing (=True) fails to export to ONXX
## Environment info
- `transformers` version: 3.5.0
- Platform: Linux-5.4.0-52-generic-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.9
- PyTorch version (GPU?): 1.7.0 (True)
- Tensorflow version (GPU?): 2.3.1 (False)
- Using GPU in script?: False
- Using distributed or parallel set-up in script?: False

### Who can help
 
Hi @LysandreJik and @patrickvonplaten, I hope I have tagged the right person. If not please untag yourself and tag the right person.

## Information

Model I am using (Bert):

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

It's a super simple script that uses gradient check-pointing with BERT.

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

A dummy data 

## To reproduce

I have made a reproducible google collab here -> https://colab.research.google.com/drive/1tUpIzbugZ4xPz6eAOJtZT-fGww9LemwN?usp=sharing

Steps to reproduce the behavior:

1. Open the notebook
2. Runtime->Run all

Error thrown:
```python
RuntimeError                              Traceback (most recent call last)

<ipython-input-5-2702a59a9c3e> in <module>()
      7         tokenizer=tokenizer,  # <-- CHANGED: add tokenizer
      8         output=Path("onnx/bert-base-cased.onnx"),
----> 9         opset=11)
     10 
     11 # Tensorflow

4 frames

/usr/local/lib/python3.6/dist-packages/torch/onnx/utils.py in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, example_outputs, opset_version, _retain_param_name, do_constant_folding, strip_doc_string, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, enable_onnx_checker, use_external_data_format, onnx_shape_inference, use_new_jit_passes)
    648                     params_dict, opset_version, dynamic_axes, defer_weight_export,
    649                     operator_export_type, strip_doc_string, val_keep_init_as_ip, custom_opsets,
--> 650                     val_add_node_names, val_use_external_data_format, model_file_location)
    651             else:
    652                 proto, export_map = graph._export_onnx(

RuntimeError: ONNX export failed: Couldn't export Python operator CheckpointFunction


```



Indeed, I see why this would fail. I don't have access to your notebook, but as a temporary workaround you could do:

```py
model.save_pretrained("here")
model = ModelClass.from_pretrained("here", gradient_checkpointing=False)
```

You should be able to convert that model to ONNX then.<|||||>I made it public now, my bad. Another work around I found was to edit the `config.json` by setting   `"gradient_checkpointing": false`. I did this because the convert script looks for the model in a path and not in memory.<|||||>Yes, this works too!<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
