Fix gather for SageMaker model parallel

huggingface · Sep 27, 2021 · 1c96500 · 1c96500
1 parent 4e0410e
commit 1c96500
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/transformers/trainer_pt_utils.py b/src/transformers/trainer_pt_utils.py
@@ -1021,6 +1021,7 @@ def smp_gather(tensor):
                 f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors."
             )
         all_tensors = smp.allgather(tensor, smp.CommGroup.DP_GROUP)
+        all_tensors = [t if len(t.shape) > 0 else t[None] for t in all_tensors]
         return torch.cat([t.cpu() for t in all_tensors], dim=0)
 
     def smp_nested_concat(tensor):