torchdata_iterable_datapipe_filter_map_pipeline_quickstart.py

1import torchdata.datapipes.iter as iter
2
3# 1. Create a source DataPipe
4data = [1, 2, 3, 4, 5, 6]
5source_dp = iter.IterableWrapper(data)
6
7# 2. Apply transformations (Filtering and Mapping)
8# Filter: only keep even numbers
9filter_dp = source_dp.filter(lambda x: x % 2 == 0)
10
11# Map: square the numbers
12map_dp = filter_dp.map(lambda x: x ** 2)
13
14# 3. Iterate through the pipeline
15for result in map_dp:
16    print(result)
17
18# Expected Output:
19# 4
20# 16
21# 36