

TITLE
Models fine-tuned with gradient checkpointing (=True) fails to export to ONXX

COMMENTS
4

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
## 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


```


