From 2ced908f815f2f13c26270c36c2027d0d9febe04 Mon Sep 17 00:00:00 2001 From: Muhammed Fatih Balin Date: Mon, 26 Aug 2024 18:13:35 -0400 Subject: [PATCH] update return type. --- python/dgl/graphbolt/sampled_subgraph.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/dgl/graphbolt/sampled_subgraph.py b/python/dgl/graphbolt/sampled_subgraph.py index 09280d37c5cb..fcb58cc572ca 100644 --- a/python/dgl/graphbolt/sampled_subgraph.py +++ b/python/dgl/graphbolt/sampled_subgraph.py @@ -255,7 +255,9 @@ def exclude_edges( ) return calling_class(*_slice_subgraph(self, index)) - def to_pyg(self, x: Union[torch.Tensor, Dict[str, torch.Tensor]]): + def to_pyg( + self, x: Union[torch.Tensor, Dict[str, torch.Tensor]] + ) -> Union[PyGLayerData, PyGLayerHeteroData]: """ Process layer inputs so that they can be consumed by a PyG model layer. @@ -266,7 +268,7 @@ def to_pyg(self, x: Union[torch.Tensor, Dict[str, torch.Tensor]]): Returns ------- - PyGLayerData + Union[PyGLayerData, PyGLayerHeteroData] A named tuple class with `x`, `edge_index` and `size` fields. Typically, a PyG GNN layer will accept these parameters. """