def collect_results_gpu(result_part, size):
    (rank, world_size) = get_dist_info()
    part_tensor = torch.tensor(bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda')
    shape_tensor = torch.tensor(part_tensor.shape, device='cuda')
    shape_list = [shape_tensor.clone() for _ in range(world_size)]
    dist.all_gather(shape_list, shape_tensor)
    shape_max = torch.tensor(shape_list).max()
    part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda')
    part_send[:shape_tensor[0]] = part_tensor
    part_recv_list = [part_tensor.new_zeros(shape_max) for _ in range(world_size)]
    dist.all_gather(part_recv_list, part_send)
    if (rank == 0):
        part_list = []
        for (recv, shape) in zip(part_recv_list, shape_list):
            part_list.append(pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()))
        ordered_results = []
        for res in zip(*part_list):
            ordered_results.extend(list(res))
        ordered_results = ordered_results[:size]
        return ordered_results
def collect_results_gpu(result_part, size):
    (rank, world_size) = get_dist_info()
    part_tensor = torch.tensor(bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda')
    shape_tensor = torch.tensor(part_tensor.shape, device='cuda')
    shape_list = [shape_tensor.clone() for _ in range(world_size)]
    dist.all_gather(shape_list, shape_tensor)
    shape_max = torch.tensor(shape_list).max()
    part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda')
    part_send[:shape_tensor[0]] = part_tensor
    part_recv_list = [part_tensor.new_zeros(shape_max) for _ in range(world_size)]
    dist.all_gather(part_recv_list, part_send)
    if (rank == 0):
        part_list = []
        for (recv, shape) in zip(part_recv_list, shape_list):
            part_list.append(pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()))
        ordered_results = []
        for res in zip(*part_list):
            ordered_results.extend(list(res))
        ordered_results = ordered_results[:size]
        return ordered_results