remove prefetcher in mixer creation

pytorch · Jan 21, 2025 · b5db9ec · b5db9ec
1 parent daafee4
commit b5db9ec
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/test/nodes/test_multi_node_weighted_sampler.py b/test/nodes/test_multi_node_weighted_sampler.py
@@ -88,12 +88,7 @@ def test_multi_node_weighted_batch_sampler_zero_weights(
 
     def test_multi_node_weighted_sampler_first_exhausted(self) -> None:
         """Test MultiNodeWeightedSampler with stop criteria FIRST_DATASET_EXHAUSTED"""
-        mixer = self._setup_multi_node_weighted_sampler(
-            self._num_samples,
-            self._num_datasets,
-            self._weights_fn,
-            stop_criteria=StopCriteria.FIRST_DATASET_EXHAUSTED,
-        )
+        mixer = MultiNodeWeightedSampler(self.datasets, self.weights, StopCriteria.FIRST_DATASET_EXHAUSTED)
 
         for _ in range(self._num_epochs):
             results = list(mixer)