Skip to content

GitLab

Unverified Commit 341dbea0 authored by Quan (Andy) Gan's avatar Quan (Andy) Gan Committed by GitHub
Browse files

[Bug] Fix unused prefetch switch for NeighborSampler and LayerSampler (#457) 

* forgot prefetch switch

* yet another stupid bug
parent 33b9c383
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
python/dgl/contrib/sampling/sampler.py
View file @ 341dbea0
...@@ -295,7 +295,7 @@ class NeighborSampler(NodeFlowSampler): ...@@ -295,7 +295,7 @@ class NeighborSampler(NodeFlowSampler):
prefetch=False, prefetch=False,
add_self_loop=False): add_self_loop=False):
super(NeighborSampler, self).__init__( super(NeighborSampler, self).__init__(
g, batch_size, seed_nodes, shuffle, num_workers * 2, g, batch_size, seed_nodes, shuffle, num_workers * 2 if prefetch else 0,
ThreadPrefetchingWrapper) ThreadPrefetchingWrapper)
assert node_prob is None, 'non-uniform node probability not supported' assert node_prob is None, 'non-uniform node probability not supported'
...@@ -361,7 +361,7 @@ class LayerSampler(NodeFlowSampler): ...@@ -361,7 +361,7 @@ class LayerSampler(NodeFlowSampler):
num_workers=1, num_workers=1,
prefetch=False): prefetch=False):
super(LayerSampler, self).__init__( super(LayerSampler, self).__init__(
g, batch_size, seed_nodes, shuffle, num_workers * 2, g, batch_size, seed_nodes, shuffle, num_workers * 2 if prefetch else 0,
ThreadPrefetchingWrapper) ThreadPrefetchingWrapper)
assert node_prob is None, 'non-uniform node probability not supported' assert node_prob is None, 'non-uniform node probability not supported'
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment