nemo_automodel.components.datasets.diffusion.loader

Typed result and construction contract for diffusion dataloaders.

Module Contents

Classes

Name	Description
`DiffusionDataloaderBuild`	Materialized diffusion dataloader, sampler, and effective batch-tail policy.
`DiffusionDataloaderConfig`	Typed construction contract used by the diffusion recipe.

API

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

Dataclass

Materialized diffusion dataloader, sampler, and effective batch-tail policy.

dataloader

DataLoader

drop_last

bool = False

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.