TFTrainer stuck in evaluation
## 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: 3.1.0
- Platform: Darwin-17.7.0-x86_64-i386-64bit
- Python version: 3.6.10
- PyTorch version (GPU?): 1.4.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: NO
- Using distributed or parallel set-up in script?: YES

### Who can help
 Trainer: @sgugger 

## Information

Model I am using (Bert, XLNet ...): TFCamembertModel and TFTrainer to fine tune model on MLM

The problem arises when using:
* [x ] the official example scripts: (give details below)
* [ ] 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
Behavior of prediction_loop function causes evaluation to go in infinite loop if prediction_loss_only is set to True. The step==steps condition is never reached
Steps to reproduce the behavior: Read code from TFTrainer below

```python
if not prediction_loss_only:
                if isinstance(logits, tuple):
                    logits = logits[0]

                if isinstance(labels, tuple):
                    labels = labels[0]

                if self.args.n_replicas > 1:
                    for val in logits.values:
                        if preds is None:
                            preds = val.numpy()
                        else:
                            preds = np.append(preds, val.numpy(), axis=0)

                    for val in labels.values:
                        if label_ids is None:
                            label_ids = val.numpy()
                        else:
                            label_ids = np.append(label_ids, val.numpy(), axis=0)
                else:
                    if preds is None:
                        preds = logits.numpy()
                    else:
                        preds = np.append(preds, logits.numpy(), axis=0)

                    if label_ids is None:
                        label_ids = labels.numpy()
                    else:
                        label_ids = np.append(label_ids, labels.numpy(), axis=0)

                if step == steps:
                    break

```

## Expected behavior
Once the number of steps reached, we should be able to get out of the for loop that draws batches

Tagging @jplu as I'm not familiar with the code of TFTrainer.<|||||>Hello!

Please provide a piece of code in order for us to be able to reproduce your issue.<|||||>Since i am using a custom dataset, here are some stats about the input:

train_dataset = tf.data.Dataset.from_tensor_slices((
    {'input_ids': masked_train_texts, 'attention_mask': tokenized_train_texts['attention_mask']},
    train_labels
))

train_dataset contains 16268 instance.

val_dataset = tf.data.Dataset.from_tensor_slices((
    {'input_ids': masked_val_texts, 'attention_mask': tokenized_val_texts['attention_mask']},
    val_labels
)) 

val_dataset contains 4067 instance.

here is only the part where i launch the trainer: 

```python
training_args = TFTrainingArguments(
    output_dir='./results',          # output directory
    num_train_epochs=4,              # total number of training epochs
    per_device_train_batch_size=64,  #64 batch size per device during training
    per_device_eval_batch_size=128,   #128 batch size for evaluation
    learning_rate=1e-5,
    evaluate_during_training=True,
    prediction_loss_only=False,
    logging_dir='./logs',            # directory for storing logs
    logging_steps=1,
    logging_first_step=True,
    eval_steps=110,
    save_steps=30,
    save_total_limit=2,
)


with training_args.strategy.scope():
    model = TFCamembertForMaskedLM.from_pretrained(CAMEMBERT_MODEL)

trainer = TFTrainer(
    model=model,                         # the instantiated 🤗 Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=train_dataset,         # training dataset
    eval_dataset=val_dataset             # evaluation dataset
)

trainer.train()
print('******')
```

After debug the TFTrainer class, i see it enters an infinite loop because if ```prediction_loss_only=True``` the condition that  stops the loop ```step==steps```<|||||>Unfortunately without a colab or a piece of code that I can copy/paste to properly reproduce your issue I cannot really help you. Nevertheless few remarks:

- Try the master version of the lib or at least the last released version (v3.4) in order to check if your problem has not been already solved.
- Which Tensorflow version you are using?
- I see in your example that you are using a masked LM model, the TF Trainer is not compliant with these models for now. The available tasks are only: Token Classification, Sequence Classification, Multiple Choice and Question Answering. The possibility to train a LM model from scratch will arrive in a future release.<|||||>I will try to reproduce this issue with some synthetic data since i cannot share my own. The problem here comes from the code. I was able to fix the issue by setting the ```prediction_loss_only``` parameter to false or by moving the ```step==steps``` condition out of the big if condition. I dont understand why it was there in the first place.

- I have checked the 3.4 version of the code and the issue is still there. It is purely an error in the code as i pointed out. I would be surprised if the code works for anyone using the same configuration as me.
- I am using Tensorflow 2.3.
-  Yes i built my tools to transform the data for a MLM tasks by following the methods implemented in Torch (with data collators).   

Globally, the TFTrainer lags very far behind the torch version functionality and quality of code wise. I am struggling with it since the beginning as it is now, whereas the torch version is much more solid and complete.  

Thank for all the great work you are doing with the lib :)<|||||>Ok, thanks, I will keep this issue open and further check this once I have a bit of time for this.<|||||>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.
