def prune_graph(adj_lists, relations):
    """The function is used to prune graph based on the relations 
    that are present in the training

    Arguments:
        adj_lists {dict} -- dictionary containing the graph
        relations {set} -- list of relation ids

    Returns:
        dict -- pruned graph
    """
    pruned_adj_list = {}
    for node, adj in adj_lists.items():
        pruned_adj_list[node] = []
        for neigh_node, rel, hp in adj:
            if rel in relations:
                pruned_adj_list[node].append((neigh_node, rel, hp))
    
    return pruned_adj_list