Fix distributed_concat with scalar tensor (#16963)

* Fix `distributed_concat` with scalar tensor

* Update trainer_pt_utils.py

Author: Yard1

src/transformers/trainer_pt_utils.py

@@ -159,8 +159,9 @@ def distributed_concat(tensor: Any, num_total_examples: Optional[int] = None) ->
     try:
         if isinstance(tensor, (tuple, list)):
             return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
+        if len(tensor.shape) <= 0:
+            tensor = tensor[None]
         output_tensors = [tensor.clone() for _ in range(dist.get_world_size())]
-        output_tensors = [t if len(t.shape) > 0 else t[None] for t in output_tensors]
         dist.all_gather(output_tensors, tensor)
         concat = torch.cat(output_tensors, dim=0)