import torch


def get_pyg_edge_index(graph, both_dir=True):
    """
    Get edge_index for an instance of torch_geometric.data.Data.

    Args:
        graph: networkx Graph
        both_dir: boolean; whether to include reverse edge from graph
    
    Return:
        torch.LongTensor of edge_index with size [2, num_edges]
    """
    sources = []
    targets = []
    for edge in graph.edges:
        sources.append(edge[0])
        targets.append(edge[1])
        if both_dir:
            sources.append(edge[1])
            targets.append(edge[0])
    return torch.LongTensor([sources, targets])