

TITLE
TFTrainer stuck in evaluation

COMMENTS
7

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

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