Repository: huggingface/transformers

Issue:
Prediction fails for certain batch sizes
## Environment info

- `transformers` version: 4.6.1
- Platform: IBM open-ce / ppcle64
- Python version: 3.8.8
- PyTorch version (GPU?): 1.7.1 NVIDIA V100
- Tensorflow version (GPU?): N/A
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: distributed (but only 1 rank per process group)

### Who can help

@sgugger


## Information

Model I am using (Bert, XLNet ...): Bert

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

Custom-implementation combining two BERT models with a multilayer perceptron that implements a regression on the last hidden layer outputs

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

a regression on pairs of protein sequence and molecule SMILES strings, plus binding affinity

## To reproduce

Steps to reproduce the behavior:

1. Call `trainer.predict()` on the encoded inputs
2. If the number of inputs is 1+`per_device_eval_batch_size`, the last batch has just one member and is treated as a scalar
3. the prediction loop fails with
``` 
File "../affinity_pred/infer_mpi.py", line 189, in main
    df_pred = predict(df)
  File "../affinity_pred/infer_mpi.py", line 180, in predict
    out = trainer.predict(x)
  File "/gpfs/alpine/world-shared/bip214/opence-env/lib/python3.8/site-packages/transformers/trainer.py", line 2065, in predict
    output = eval_loop(
  File "/gpfs/alpine/world-shared/bip214/opence-env/lib/python3.8/site-packages/transformers/trainer.py", line 2162, in evaluation_loop
    logits = self._nested_gather(logits)
  File "/gpfs/alpine/world-shared/bip214/opence-env/lib/python3.8/site-packages/transformers/trainer.py", line 2252, in _nested_gather
    tensors = distributed_concat(tensors)
  File "/gpfs/alpine/world-shared/bip214/opence-env/lib/python3.8/site-packages/transformers/trainer_pt_utils.py", line 158, in distributed_concat
    concat = torch.cat(output_tensors, dim=0)
RuntimeError: zero-dimensional tensor (at position 0) cannot be concatenated
```

**Example**:

last two batches (`per_device_eval_batch_size` is 32)
```
[tensor([ 0.5737, -0.1372, -0.5283, -0.0641, -0.0641,  0.6353,  0.1035, -0.4148,
         0.2314, -0.3879, -0.4431, -0.3931, -0.2642,  0.5039, -0.4187,  0.0679,
         0.0679, -0.3167, -0.4783, -0.6724, -0.6724, -0.3257,  0.4922,  0.4922,
        -0.4189, -0.3652, -0.4468, -0.2358, -0.3696,  0.1646, -0.2004, -1.0234],
       device='cuda:0', dtype=torch.float16)]
[tensor(0.7144, device='cuda:0', dtype=torch.float16)]
RuntimeError('zero-dimensional tensor (at position 0) cannot be concatenated')
```

## Expected behavior

No exception


Discussion:

Resolving patch:
diff --git a/src/transformers/trainer_pt_utils.py b/src/transformers/trainer_pt_utils.py
--- a/src/transformers/trainer_pt_utils.py
+++ b/src/transformers/trainer_pt_utils.py
@@ -155,6 +155,7 @@ def distributed_concat(tensor: "torch.Tensor", num_total_examples: Optional[int]
             return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
         output_tensors = [tensor.clone() for _ in range(dist.get_world_size())]
         dist.all_gather(output_tensors, tensor)
+        output_tensors = [t if len(t.shape) > 0 else t[None] for t in output_tensors]
         concat = torch.cat(output_tensors, dim=0)
 
         # truncate the dummy elements added by SequentialDistributedSampler


Test patch: