Function gather_tensor

core/mengine.py:25–28

(tensor, n)

def gather_tensor(tensor, n):
    rt = [torch.zeros_like(tensor) for _ in range(n)]
    dist.all_gather(rt, tensor)
    return torch.cat(rt, dim=0)


class TrainEngine(object):