Single-Batch Data Loader

Module for the Single-Batch Data Loader class.

class _SingleBatchDataLoader(dataset)

Bases: object

Data loader wrapper that returns the entire dataset as a single batch.

This utility is intended for cases where mini-batching is disabled (e.g. batch_size=None). The loader yields exactly one batch per iteration.

In distributed environments, the dataset is automatically partitioned across processes according to the current rank and world size. Each process receives only its corresponding subset of data.

In non-distributed environments, the full dataset is returned.

Initialization of the _SingleBatchDataLoader class.

In distributed training, the dataset indices are split across processes using the current rank and world size, so that each process receives only its corresponding subset of data.

In non-distributed training, the full dataset is loaded.

The resulting data is converted into a single batch and stored internally.

Parameters:

dataset – Dataset object.

Raises:

RuntimeError – If the dataset size is smaller than the number of distributed processes.