update return type.

dmlc · Aug 26, 2024 · 2ced908 · 2ced908
1 parent f37a455
commit 2ced908
Showing 1 changed file with 4 additions and 2 deletions.
diff --git 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.
         """