blabal

dmlc · Aug 2, 2024 · 2e4956c · 2e4956c
1 parent 1f1316d
commit 2e4956c
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 5 deletions.
diff --git a/python/dgl/graphbolt/dataloader.py b/python/dgl/graphbolt/dataloader.py
@@ -4,14 +4,18 @@
 
 import torch
 import torch.utils.data as torch_data
-from torch.utils.data.graph import traverse_dps
 
 from .base import CopyTo, get_host_to_device_uva_stream
 from .feature_fetcher import FeatureFetcher, FeatureFetcherStartMarker
 from .impl.gpu_graph_cache import GPUGraphCache
 from .impl.neighbor_sampler import SamplePerLayer
 
-from .internal import datapipe_graph_to_adjlist, find_dps, replace_dp
+from .internal import (
+    datapipe_graph_to_adjlist,
+    find_dps,
+    replace_dp,
+    traverse_dps,
+)
 from .item_sampler import ItemSampler
 
 

diff --git a/python/dgl/graphbolt/internal/datapipe_utils.py b/python/dgl/graphbolt/internal/datapipe_utils.py
@@ -5,7 +5,12 @@
 from torch.utils.data import IterDataPipe, MapDataPipe
 from torch.utils.data.graph import DataPipe, DataPipeGraph, traverse_dps
 
-__all__ = ["datapipe_graph_to_adjlist", "find_dps", "replace_dp"]
+__all__ = [
+    "datapipe_graph_to_adjlist",
+    "find_dps",
+    "replace_dp",
+    "traverse_dps",
+]
 
 
 def _get_parents(result_dict, datapipe_graph):

diff --git a/tests/python/pytorch/graphbolt/test_dataloader.py b/tests/python/pytorch/graphbolt/test_dataloader.py
@@ -7,8 +7,7 @@
 import pytest
 import torch
 
-from dgl.graphbolt.internal import find_dps
-from torch.utils.data.graph import traverse_dps
+from dgl.graphbolt.internal import find_dps, traverse_dps
 
 from . import gb_test_utils