Strange ValueError with GPT-2
## 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.4.2
- Platform: macOS-10.15.7-x86_64-i386-64bit
- Python version: 3.8.6
- PyTorch version (GPU?): 1.7.1 (False)
- Tensorflow version (GPU?): 2.4.0 (False)
- Using GPU in script?: No.
- Using distributed or parallel set-up in script?: No.

### Who can help

Models:

- gpt2: @patrickvonplaten, @LysandreJik

## Information

Model I am using (Bert, XLNet ...): GPT-2

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

Create and train GPT-2 with 8 heads and batch size 32.

I get an error message when training GPT-2. When number of heads and batch size are the same it works. Looks like a shape check is wronk. See error message below.

```
  File "train.py", line 66, in <module>
    history = model.fit(
  File "/Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py", line 1100, in fit
    tmp_logs = self.train_function(iterator)
  File "/Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
  File "/Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 871, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 725, in _initialize
    self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
  File "/Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 2969, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3361, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3196, in _create_graph_function
    func_graph_module.func_graph_from_py_func(
  File "/Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 977, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:805 train_function  *
        return step_function(self, iterator)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:795 step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:1259 run
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2730 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:3417 _call_for_each_replica
        return fn(*args, **kwargs)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:788 run_step  **
        outputs = model.train_step(data)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:758 train_step
        self.compiled_metrics.update_state(y, y_pred, sample_weight)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/keras/engine/compile_utils.py:408 update_state
        metric_obj.update_state(y_t, y_p, sample_weight=mask)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/keras/utils/metrics_utils.py:90 decorated
        update_op = update_state_fn(*args, **kwargs)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/keras/metrics.py:177 update_state_fn
        return ag_update_state(*args, **kwargs)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/keras/metrics.py:618 update_state  **
        matches = ag_fn(y_true, y_pred, **self._fn_kwargs)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper
        return target(*args, **kwargs)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/keras/metrics.py:3315 sparse_categorical_accuracy
        return math_ops.cast(math_ops.equal(y_true, y_pred), K.floatx())
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py:201 wrapper
        return target(*args, **kwargs)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py:1679 equal
        return gen_math_ops.equal(x, y, name=name)
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py:3177 equal
        _, _, _op, _outputs = _op_def_library._apply_op_helper(
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/framework/op_def_library.py:748 _apply_op_helper
        op = g._create_op_internal(op_type_name, inputs, dtypes=None,
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py:590 _create_op_internal
        return super(FuncGraph, self)._create_op_internal(  # pylint: disable=protected-access
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/framework/ops.py:3528 _create_op_internal
        ret = Operation(
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/framework/ops.py:2015 __init__
        self._c_op = _create_c_op(self._graph, node_def, inputs,
    /Users/tristanbehrens/Development/python-venvs/tf2-p38/lib/python3.8/site-packages/tensorflow/python/framework/ops.py:1856 _create_c_op
        raise ValueError(str(e))

    ValueError: Dimensions must be equal, but are 32 and 8 for '{{node Equal_1}} = Equal[T=DT_FLOAT, incompatible_shape_error=true](Cast_6, Cast_7)' with input shapes: [32,301], [2,32,8,301].

```

Hi, could you give a reproducible code example?<|||||>With pleasure!

```
import tensorflow as tf
from transformers import GPT2Config, TFGPT2LMHeadModel
import numpy as np

# Prepare the dataset.
inputs = np.random.randint(0, 128, (1025, 256))
labels = np.random.randint(0, 128, (1025, 256))
dataset = tf.data.Dataset.from_tensor_slices((inputs, labels))
dataset = dataset.shuffle(1000).batch(32, drop_remainder=True)

# Create the model configuration.
model_config = GPT2Config(
    vocab_size=128,
    bos_token_id=0,
    eos_token_id=127,
    n_head=8,
    n_layer=6,
    n_embd=512,
    n_positions=2048
)
print(model_config)

# Create the model.
model = TFGPT2LMHeadModel(model_config)

# Compile the model.
optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
model.compile(
    optimizer=optimizer,
    loss=[loss, *[None] * model.config.n_layer],
    metrics=[metric]
)

# Train the model.
history = model.fit(
    dataset, 
    epochs=10,
)

```<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.