Update _utils.py to fix merge batch

hpcaitech · Aug 28, 2023 · 5f4984f · 5f4984f
1 parent 80e9e96
commit 5f4984f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/colossalai/pipeline/schedule/_utils.py b/colossalai/pipeline/schedule/_utils.py
@@ -119,14 +119,14 @@ def merge_batch(data: List[Any]) -> Any:
     tree_spec = None
     for d in data:
         elems, tree_spec = tree_flatten(d)
-        flattened_data.append(elems)
+        flattened_data.append(elems[0])
     merged_data = []
-    for elem_batch in zip(*flattened_data):
+    for elem_batch in zip(*(fd.items() for fd in flattened_data)):
         if isinstance(elem_batch[0], torch.Tensor):
             if len(elem_batch[0].shape) == 0:    # set loss to None in pipeline outputs
                 merged_data.append(None)
             else:
                 merged_data.append(torch.cat(elem_batch, dim=0))
         else:
             merged_data.append(list(elem_batch))
-    return tree_unflatten(merged_data, tree_spec)
+    return tree_unflatten([merged_data], tree_spec)