[Bug] Fix unused prefetch switch for NeighborSampler and LayerSampler (…

…dmlc#457) * forgot prefetch switch * yet another stupid bug
liuweiping2020 · Mar 19, 2019 · 341dbea · 341dbea
1 parent 33b9c38
commit 341dbea
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/python/dgl/contrib/sampling/sampler.py b/python/dgl/contrib/sampling/sampler.py
@@ -165,7 +165,7 @@ def __init__(
 
         if num_prefetch:
             self._prefetching_wrapper_class = prefetching_wrapper_class
-            self._num_prefetch = num_prefetch
+        self._num_prefetch = num_prefetch
 
     def fetch(self, current_nodeflow_index):
         '''
@@ -295,7 +295,7 @@ def __init__(
             prefetch=False,
             add_self_loop=False):
         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)
 
         assert node_prob is None, 'non-uniform node probability not supported'
@@ -361,7 +361,7 @@ def __init__(
             num_workers=1,
             prefetch=False):
         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)
 
         assert node_prob is None, 'non-uniform node probability not supported'