nemo_automodel.components.datasets.diffusion.loader

Typed result and construction contract for diffusion dataloaders.

Module Contents

Classes

Name	Description
`DiffusionDataloaderBuild`	Materialized diffusion dataloader and its sampler.
`DiffusionDataloaderConfig`	Typed construction contract used by the diffusion recipe.

Data

__all__

API

class nemo_automodel.components.datasets.diffusion.loader.DiffusionDataloaderBuild(
    dataloader: torch.utils.data.DataLoader,
    sampler: torch.utils.data.Sampler[object] | None
)

Dataclass

Materialized diffusion dataloader and its sampler.

dataloader

DataLoader

sampler

Sampler[object] | None

class nemo_automodel.components.datasets.diffusion.loader.DiffusionDataloaderConfig()

Protocol

Typed construction contract used by the diffusion recipe.

nemo_automodel.components.datasets.diffusion.loader.DiffusionDataloaderConfig.build(
    dp_rank: int,
    dp_world_size: int,
    batch_size: int
) -> nemo_automodel.components.datasets.diffusion.loader.DiffusionDataloaderBuild

Build the configured per-rank diffusion dataloader.

nemo_automodel.components.datasets.diffusion.loader.__all__ = ['DiffusionDataloaderBuild', 'DiffusionDataloaderConfig']