def test_hetero_graph_batch(self):
G = generate_simple_hete_graph()
hete = HeteroGraph(G)
hete = HeteroGraph(
node_feature=hete.node_feature,
node_label=hete.node_label,
edge_feature=hete.edge_feature,
edge_label=hete.edge_label,
edge_index=hete.edge_index,
directed=True
)
heteGraphDataset = []
for _ in range(30):
heteGraphDataset.append(hete.clone())
dataloader = DataLoader(
heteGraphDataset,
collate_fn=Batch.collate(),
batch_size=3,
shuffle=True,
)
self.assertEqual(len(dataloader), math.ceil(30 / 3))
for data in dataloader:
self.assertEqual(data.num_graphs, 3)
def test_hetero_graph_batch(self):
G = generate_simple_hete_graph()
hete = HeteroGraph(G)
heteGraphDataset = []
for i in range(30):
heteGraphDataset.append(hete.clone())
dataloader = DataLoader(heteGraphDataset,
collate_fn=Batch.collate(),
batch_size=3,
shuffle=True)
self.assertEqual(len(dataloader), math.ceil(30 / 3))
for data in dataloader:
self.assertEqual(data.num_graphs, 3)