Method dataset

axlearn/common/input_base.py:191–201

Returns the input dataset, which should produce per-feed logical batches. Each batch is a pytree of arrays which reside on host memory (i.e., leaves can be any array type which can be converted to numpy via `as_numpy_array`). The dataset should be iterable, i.e., it is expe

(self)

        )

    def dataset(self) -> Iterable[Nested[Tensor]]:
        """Returns the input dataset, which should produce per-feed logical batches.

        Each batch is a pytree of arrays which reside on host memory (i.e., leaves can be any array
        type which can be converted to numpy via `as_numpy_array`).

        The dataset should be iterable, i.e., it is expected to support conversion to an iterator
        via `iter(...)`. Although not strictly required, it is recommended for the iterator to be
        checkpointable.
        """
        raise NotImplementedError(type(self))

    def __iter__(self) -> Iterator[Nested[Tensor]]:
        """Iterates over the input dataset.