fix imports

rapidsai · Aug 22, 2024 · 1fcd982 · 1fcd982
1 parent d8916c6
commit 1fcd982
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/python/cugraph/cugraph/gnn/data_loading/__init__.py b/python/cugraph/cugraph/gnn/data_loading/__init__.py
@@ -14,9 +14,11 @@
 from cugraph.gnn.data_loading.bulk_sampler import BulkSampler
 from cugraph.gnn.data_loading.dist_sampler import (
     DistSampler,
+    NeighborSampler,
+)
+from cugraph.gnn.data_loading.dist_io import (
     DistSampleWriter,
     DistSampleReader,
-    NeighborSampler,
 )