Use validator from torch main

Author: bushshrub

torchdata/datapipes/iter/transform/callable.py

@@ -8,7 +8,7 @@
 from typing import Callable, Hashable, Iterator, List, Optional, Set, Sized, TypeVar, Union
 
 from torch.utils.data import functional_datapipe, IterDataPipe
-from torch.utils.data.datapipes.utils.common import _check_unpickable_fn
+from torch.utils.data.datapipes.utils.common import _check_unpickable_fn, validate_input_col
 
 T_co = TypeVar("T_co", covariant=True)
 
@@ -103,6 +103,10 @@ class FlatMapperIterDataPipe(IterDataPipe[T_co]):
     Args:
         datapipe: Source IterDataPipe
         fn: the function to be applied to each element in the DataPipe, the output must be a Sequence
+        input_col: Index or indices of data which ``fn`` is applied, such as:
+            - ``None`` as default to apply ``fn`` to the data directly.
+            - Integer(s) is/are used for list/tuple.
+            - Key(s) is/are used for dict.
 
     Example:
         >>> from torchdata.datapipes.iter import IterableWrapper
@@ -122,6 +126,7 @@ def __init__(self, datapipe: IterDataPipe, fn: Callable, input_col=None) -> None
         _check_unpickable_fn(fn)
         self.fn = fn  # type: ignore[assignment]
         self.input_col = input_col
+        validate_input_col(fn, input_col)
 
     def _apply_fn(self, data):
         if self.input_col is None: