import dgl
import torch as th
from dgl.data.utils import save_graphs
g1 = dgl.DGLGraph()
g1.add_nodes(3)
g1.add_edges([0, 0, 0, 1, 1, 2], [0, 1, 2, 1, 2, 2])
g1.ndata["x"] = th.ones(3, 5)
g1.edata['y'] = th.zeros(6, 5)
g2 = dgl.DGLGraph()
g2.add_nodes(3)
g2.add_edges([0, 1, 2], [1, 2, 1])
g2.edata["e"] = th.ones(3, 4)
graph_labels = {"graph_sizes": th.tensor([3, 3])}
save_graphs("data/try1.bin", [g1, g2], graph_labels)
from dgl.data.utils import load_graphs
from dgl.data.utils import load_labels
glist, label_dict = load_graphs("data/try1.bin", [0])
graph_sizes = load_labels("data/try1.bin")
print(glist)
print(label_dict)
print(graph_sizes)