add SamplerWrapper to init (#1360)

pytorch · Nov 13, 2024 · db03884 · db03884
1 parent 7fa74e4
commit db03884
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchdata/nodes/__init__.py b/torchdata/nodes/__init__.py
@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .adapters import IterableWrapper, MapStyleWrapper
+from .adapters import IterableWrapper, MapStyleWrapper, SamplerWrapper
 from .base_node import BaseNode, T
 from .batch import Batcher
 from .loader import Loader
@@ -25,6 +25,7 @@
     "ParallelMapper",
     "PinMemory",
     "Prefetcher",
+    "SamplerWrapper",
     "Stateful",
     "T",
 ]