in dualpipe/utils.py [0:0]
def gather(micro_outputs, dim):
assert isinstance(micro_outputs[0], (torch.Tensor, tuple, list))
if isinstance(micro_outputs[0], torch.Tensor):
micro_outputs = [(x,) for x in micro_outputs]
outputs = [x for x in zip(*micro_outputs)]
outputs = tuple(cat_tensor(x, dim=dim) for x in outputs)
return outputs