Model.fit on GPT2 and TPUs
## Environment info
- `transformers` version: 3.0.2
- Platform: Linux-4.19.112+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.6.0+cu101 (False)
- Tensorflow version (GPU?): 2.3.0 (False)
- Using GPU in script?:  No
- Using distributed or parallel set-up in script?: Yes

## Information

Model I am using: GPT2

@jplu When using the keras model.fit() method, it looks like there's a problem with logits and tensors: `Compilation failure: logits and labels must have the same first dimension`. Setting `from_logits = False` doesn't seem to resolve the problem. Any suggestion on how to change model compilation or the dataset to fix this?

```
# TPU and Strategy Initialization
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)

# Load and compile model
with strategy.scope():
    model = TFGPT2LMHeadModel.from_pretrained('gpt2-medium')
    model.resize_token_embeddings(len(tokenizer))
    optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5)
    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    model.compile(optimizer = optimizer,
              loss = loss)
 
# Format inputs
input_ids = tf.convert_to_tensor(input_ids)
attention_mask = tf.convert_to_tensor(attention_mask)
labels = tf.convert_to_tensor(labels)

# Create dataset
dataset = tf.data.Dataset.from_tensor_slices(({'input_ids': input_ids, 'attention_mask': attention_mask}, labels))
num_examples = tf.data.experimental.cardinality(dataset).numpy()
train_dataset = dataset.repeat().shuffle(num_examples).batch(8)

# Train model
model.fit(x = train_dataset, epochs = 1, batch_size = 8, steps_per_epoch = num_examples)
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-10-4efb969ef56e> in <module>()
---> 23 model.fit(x = train_dataset, epochs = 1, batch_size = 8, steps_per_epoch = num_examples)

InvalidArgumentError: 9 root error(s) found.
  (0) Invalid argument: {{function_node __inference_train_function_121591}} Compilation failure: logits and labels must have the same first dimension, got logits shape [32768,64] and labels shape [1024]
	 [[{{node sparse_categorical_crossentropy_24/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits}}]]
	TPU compilation failed
	 [[tpu_compile_succeeded_assert/_142128106034484350/_6]]
	 [[tpu_compile_succeeded_assert/_142128106034484350/_6/_223]]
  (1) Invalid argument: {{function_node __inference_train_function_121591}} Compilation failure: logits and labels must have the same first dimension, got logits shape [32768,64] and labels shape [1024]
	 [[{{node sparse_categorical_crossentropy_24/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits}}]]
	TPU compilation failed
	 [[tpu_compile_succeeded_assert/_142128106034484350/_6]]
	 [[tpu_compile_succeeded_assert/_142128106034484350/_6/_307]]
```

Humm for me it looks like it is an issue with the dataset creation, but I might be wrong as I don't have the code that creates the features.

Can you try without the `steps_per_epoch` parameter?<|||||>isn't steps_per_epoch equal to `num_examples/batch_size` ? futhermore your labels first  dimention (batch) and dataset batch dimention  must be equal i.e #  rows of dataset/X == # rows of labels `where 32768 !=1024`<|||||>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.
