def split_train_test_graph(graph):
    """
    For test true edges, 1/3 of the edges are randomly chosen
    and removed as ground truth in the test set,
    the remaining graph is taken as the training set.
    """
    test_edges = []
    neg_sampler = dgl.dataloading.negative_sampler.Uniform(1)
    sampled_edge_ids = random.sample(
        range(graph.num_edges()), int(graph.num_edges() / 3)
    )
    for edge_id in sampled_edge_ids:
        src, dst = graph.find_edges(edge_id)
        test_edges.append(TestEdge(src, dst, 1))

        src, dst = neg_sampler(graph, th.tensor([edge_id]))
        test_edges.append(TestEdge(src, dst, 0))

    graph.remove_edges(sampled_edge_ids)
    test_graph = test_edges

    return graph, test_graph

    For test true edges, 1/3 of the edges are randomly chosen
    and removed as ground truth in the test set,
    the remaining graph is taken as the training set.
    
examples/pytorch/eges/utils.py