add SamplerWrapper to init (pytorch#1360)

andrewkho · web-flow · commit db0388430438 · 2024-11-13T13:05:08.000-08:00
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",
 ]
