Skip to navigation
DevelopmentAPI ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsDatasetsDiffusionnemo_automodel.components.datasets.diffusion.loader

nemo_automodel.components.datasets.diffusion.loader

View as Markdown

Typed result and construction contract for diffusion dataloaders.

Module Contents

Classes

NameDescription
DiffusionDataloaderBuildMaterialized diffusion dataloader, sampler, and effective batch-tail policy.
DiffusionDataloaderConfigTyped 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.