[BUG] Trainer predict bug under DDP model.
Author: hijkzzz
Pull request: no

### Background

The model is trained with DDP.
The error denotes that the batch is smaller than the model required.
However, the test file cannot use `droplast`.
How I can predict the test file with DDP or remove DDP? 

### Code
```
training_args = TrainingArguments(
    output_dir='./results',          # output directory
    num_train_epochs=10,              # total number of training epochs
    per_device_train_batch_size=2,  # batch size per device during training
    per_device_eval_batch_size=8,   # batch size for evaluation
    warmup_ratio=0.1,                # number of warmup steps for learning rate scheduler
    weight_decay=1e-2,               # strength of weight decay
    logging_dir='./logs',            # directory for storing logs
    logging_steps=100,
    evaluation_strategy='epoch',
    gradient_accumulation_steps=4,
    metric_for_best_model="f1",
    fp16=True
)

trainer = Trainer(
    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
    compute_metrics=compute_metrics
)

trainer.train()

test_dataset = WNUTDataset(test_encodings)

predictions, labels, _ = trainer.predict(test_dataset)
predictions = np.argmax(predictions, axis=2)
```

### Bugs
```
~/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/_functions.py in forward(ctx, target_device, dim, *inputs)
     66             ctx.unsqueezed_scalar = False
     67         ctx.input_sizes = tuple(map(lambda i: i.size(ctx.dim), inputs))
---> 68         return comm.gather(inputs, ctx.dim, ctx.target_device)
     69 
     70     @staticmethod

~/anaconda3/lib/python3.7/site-packages/torch/cuda/comm.py in gather(tensors, dim, destination)
    164         concatenating ``tensors`` along ``dim``.
    165     """
--> 166     return torch._C._gather(tensors, dim, destination)

RuntimeError: Gather got an input of invalid size: got [1024, 7, 768], but expected [1024, 8, 768]
```

@sgugger

Comments:
Without seeing the whole stack trace, your version of Transformers used (please follow the issue template!) or the code you are using to build your dataset,  there is little we can do to help.

> Without seeing the whole stack trace, your version of Transformers used (please follow the issue template!) or the code you are using to build your dataset, there is little we can do to help.

Hi, the code is in https://colab.research.google.com/drive/1wmdjmU54iVSGXeJXzLO706uRzg2hlVZb?usp=sharing

At present, **I change the evaluate batch size to 1, and the prediction is successful. But it' s very slow.**

Note that I trained the model offline, not in the colab. I think maybe `Transformers` should provide a api to specify parallel traning model (the defaut is nn.DataParallel, however ... the bugs).

```
# load the best model, batch_size = 1 (for DDP bug, batch_size=8 get a error)
training_args2 = TrainingArguments(
    output_dir='./results',          # output directory
    per_device_eval_batch_size=1,   # batch size for evaluation
    logging_dir='./logs',            # directory for storing logs
    logging_steps=100,
)

trainer2 = Trainer(
    model=trainer.model,                         # the instantiated   Transformers model to be trained
    args=training_args2,                  # training arguments, defined above
    compute_metrics=compute_metrics
)
```



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.