Did the underlying pre-trained models change somehow?
Author: jmwoloso
Pull request: no

## 🐛 Bug

<!-- Important information -->

Model I am using (Bert, XLNet....): Bert

Language I am using the model on (English, Chinese....): English

The problem arise when using:
* [ ] the official example scripts: (give details)
* [X] my own modified scripts: I'm just loading the pre-trained bert-base-uncased model

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: document classification

## To Reproduce

Steps to reproduce the behavior:

1. `PRETRAINED_WEIGHTS = "bert-base-uncased"`
2. `model = TFBertForSequenceClassification.from_pretrained(PRETRAINED_WEIGHTS)`
3. see just below
```
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
sparse_categorical_accuracy = tf.keras.metrics.SparseCategoricalAccuracy("train_accuracy")
model.compile(optimizer=optimizer, 
              loss=loss,
              metrics=[sparse_categorical_accuracy])
```
4. see just below
```
history = model.fit([train_input_ids, train_input_masks, train_segment_ids],
                    train_labels,
                    epochs=EPOCHS,
                    batch_size=BATCH_SIZE,
                    validation_data=([val_input_ids, val_input_masks, val_segment_ids],
                                     val_labels),
                    use_multiprocessing=True,
                    verbose=1)
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well. -->

## Expected behavior

That the model begins training.

## Environment

* OS: Databricks (linux?)
* Python version: 3.7
* PyTorch version: I'm using the TF-flavor models
* PyTorch Transformers version (or branch): v2.1.1 & v2.2.0 (see additional context)
* Using GPU ? Yes
* Distributed of parallel setup ? No
* Any other relevant information:

## Additional context
In databricks I was not pinning the transformers version so it was installing the latest. I have always pinned the `mlflow` version to 1.4.0 though. I tried with what the latest prior release of `transformers` was (2.1.1) and still got the same error where this worked flawlessly before. The error is below and it specifies it is an `mlflow` issue, though in reality I think it may have something to do with the pretrained model that is loaded when we specify `bert-base-uncased`. It seems this underlying model changed independently of the latest release of `transformers`? Are the pre-trained models from some public Google repository or are they Huggingface-specific?

Thanks again for supporting TF 2, this repo has been a blessing!

Traceback:

```
UserWarning: Logging to MLflow failed: Saving the model to HDF5 format requires the model to be a Functional model or a Sequential model. It does not work for subclassed models, because such models are defined via the body of a Python method, which isn't safely serializable. Consider saving to the Tensorflow SavedModel format (by setting save_format="tf") or using `save_weights`.
  try_mlflow_log(mlflow.keras.log_model, self.model, artifact_path='model')



InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  Incompatible shapes: [4,12,512,512] vs. [4,1,1,0]
	 [[node tf_bert_for_sequence_classification/bert/encoder/layer_._0/attention/self/add (defined at /databricks/python/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1751) ]]
	 [[cluster_1_1/xla_compile]]
	 [[cluster_0_1/merge_oidx_1/_22]]
  (1) Invalid argument:  Incompatible shapes: [4,12,512,512] vs. [4,1,1,0]
	 [[node tf_bert_for_sequence_classification/bert/encoder/layer_._0/attention/self/add (defined at /databricks/python/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1751) ]]
	 [[cluster_1_1/xla_compile]]
0 successful operations.
0 derived errors ignored. [Op:__inference_distributed_function_37652]

Function call stack:
distributed_function -> distributed_function
```
<!-- Add any other context about the problem here. -->


Comments:
Total user-error here. Made a change to a function that writes the TF Records to storage (the name of one of the features) and didn't propagate that info to the function i wrote that reads the TF Records back in, so it wasn't loading my input masks because it was looking for a key that didn't exist.