> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.datasets.diffusion.mock_dataloader

Mock dataloader for automodel WAN training tests.

This module provides a mock dataset and dataloader that generates random
tensors with the correct shapes for WAN 2.1 training, allowing functional
tests to run without requiring real data.

## Module Contents

### Classes

| Name                                                                                                               | Description                                                               |
| ------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------- |
| [`MockWanDataloaderConfig`](#nemo_automodel-components-datasets-diffusion-mock_dataloader-MockWanDataloaderConfig) | Construction-time configuration for a mock WAN dataloader.                |
| [`MockWanDataset`](#nemo_automodel-components-datasets-diffusion-mock_dataloader-MockWanDataset)                   | Mock dataset that generates random data matching WAN 2.1 expected format. |
| [`MockWanDatasetConfig`](#nemo_automodel-components-datasets-diffusion-mock_dataloader-MockWanDatasetConfig)       | Construction-time configuration for :class:`MockWanDataset`.              |

### Functions

| Name                                                                                                           | Description                                                                |
| -------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`build_mock_dataloader`](#nemo_automodel-components-datasets-diffusion-mock_dataloader-build_mock_dataloader) | Build a mock dataloader for WAN training tests.                            |
| [`mock_collate_fn`](#nemo_automodel-components-datasets-diffusion-mock_dataloader-mock_collate_fn)             | Collate function for mock dataset, matching the real collate\_fn behavior. |

### API

```python
class nemo_automodel.components.datasets.diffusion.mock_dataloader.MockWanDataloaderConfig(
    num_workers: int = 0,
    device: str = 'cpu',
    length: int = 1024,
    num_channels: int = 16,
    num_frame_latents: int = 16,
    spatial_h: int = 30,
    spatial_w: int = 52,
    text_seq_len: int = 77,
    text_embed_dim: int = 4096,
    shuffle: bool = True
)
```

Dataclass

Construction-time configuration for a mock WAN dataloader.

**`device`** `str = 'cpu'`

---

**`length`** `int = 1024`

---

**`num_channels`** `int = 16`

---

**`num_frame_latents`** `int = 16`

---

**`num_workers`** `int = 0`

---

**`shuffle`** `bool = True`

---

**`spatial_h`** `int = 30`

---

**`spatial_w`** `int = 52`

---

**`text_embed_dim`** `int = 4096`

---

**`text_seq_len`** `int = 77`

---

```python
nemo_automodel.components.datasets.diffusion.mock_dataloader.MockWanDataloaderConfig.build(
    dp_rank: int,
    dp_world_size: int,
    batch_size: int
) -> nemo_automodel.components.datasets.diffusion.loader.DiffusionDataloaderBuild
```

Build the configured mock dataset, sampler, and dataloader.

```python
class nemo_automodel.components.datasets.diffusion.mock_dataloader.MockWanDataset(
    length: int = 1024,
    num_channels: int = 16,
    num_frame_latents: int = 16,
    spatial_h: int = 30,
    spatial_w: int = 52,
    text_seq_len: int = 77,
    text_embed_dim: int = 4096,
    device: str = 'cpu'
)
```

**Bases:** `Dataset`

Mock dataset that generates random data matching WAN 2.1 expected format.

**Parameters:**

**`length`** `int` — default: 1024

Number of samples in the dataset.

---

**`num_channels`** `int` — default: 16

Number of latent channels (default: 16 for WAN).

---

**`num_frame_latents`** `int` — default: 16

Number of temporal latent frames.

---

**`spatial_h`** `int` — default: 30

Height of spatial latents.

---

**`spatial_w`** `int` — default: 52

Width of spatial latents.

---

**`text_seq_len`** `int` — default: 77

Length of text sequence.

---

**`text_embed_dim`** `int` — default: 4096

Dimension of text embeddings (default: 4096 for UMT5).

---

**`device`** `str` — default: 'cpu'

Device to place tensors on.

---

**`length`** `= max(int(length), 1)`

---

```python
nemo_automodel.components.datasets.diffusion.mock_dataloader.MockWanDataset.__getitem__(
    idx: int
) -> typing.Dict[str, torch.Tensor]
```

Generate a mock sample with random data.

**Returns:** `Dict[str, torch.Tensor]`

Dict containing:

* text\_embeddings: \[1, text\_seq\_len, text\_embed\_dim]
* video\_latents: \[1, num\_channels, num\_frame\_latents, spatial\_h, spatial\_w]
* metadata: empty dict
* file\_info: mock file info

```python
nemo_automodel.components.datasets.diffusion.mock_dataloader.MockWanDataset.__len__() -> int
```

```python
class nemo_automodel.components.datasets.diffusion.mock_dataloader.MockWanDatasetConfig(
    length: int = 1024,
    num_channels: int = 16,
    num_frame_latents: int = 16,
    spatial_h: int = 30,
    spatial_w: int = 52,
    text_seq_len: int = 77,
    text_embed_dim: int = 4096,
    device: str = 'cpu'
)
```

Dataclass

Construction-time configuration for :class:`MockWanDataset`.

**`device`** `str = 'cpu'`

Device to place tensors on.

---

**`length`** `int = 1024`

Number of samples in the dataset.

---

**`num_channels`** `int = 16`

Number of latent channels.

---

**`num_frame_latents`** `int = 16`

Number of temporal latent frames.

---

**`spatial_h`** `int = 30`

Height of spatial latents.

---

**`spatial_w`** `int = 52`

Width of spatial latents.

---

**`text_embed_dim`** `int = 4096`

Dimension of text embeddings.

---

**`text_seq_len`** `int = 77`

Length of text sequence.

---

```python
nemo_automodel.components.datasets.diffusion.mock_dataloader.MockWanDatasetConfig.build() -> 'MockWanDataset'
```

Build a :class:`MockWanDataset` from this :class:`MockWanDatasetConfig`.

```python
nemo_automodel.components.datasets.diffusion.mock_dataloader.build_mock_dataloader(
    dp_rank: int = 0,
    dp_world_size: int = 1,
    batch_size: int = 1,
    num_workers: int = 0,
    device: str = 'cpu',
    length: int = 1024,
    num_channels: int = 16,
    num_frame_latents: int = 16,
    spatial_h: int = 30,
    spatial_w: int = 52,
    text_seq_len: int = 77,
    text_embed_dim: int = 4096,
    shuffle: bool = True
) -> typing.Tuple[torch.utils.data.DataLoader, torch.utils.data.DistributedSampler | None]
```

Build a mock dataloader for WAN training tests.

This function follows the same interface as build\_dataloader but generates
random data instead of loading from .meta files.

**Parameters:**

**`dp_rank`** `int` — default: 0

Data parallel rank.

---

**`dp_world_size`** `int` — default: 1

Data parallel world size.

---

**`batch_size`** `int` — default: 1

Batch size per GPU.

---

**`num_workers`** `int` — default: 0

Number of dataloader workers.

---

**`device`** `str` — default: 'cpu'

Device to place tensors on.

---

**`length`** `int` — default: 1024

Number of samples in mock dataset.

---

**`num_channels`** `int` — default: 16

Number of latent channels (default: 16).

---

**`num_frame_latents`** `int` — default: 16

Number of temporal latent frames.

---

**`spatial_h`** `int` — default: 30

Height of spatial latents.

---

**`spatial_w`** `int` — default: 52

Width of spatial latents.

---

**`text_seq_len`** `int` — default: 77

Length of text sequence.

---

**`text_embed_dim`** `int` — default: 4096

Dimension of text embeddings.

---

**`shuffle`** `bool` — default: True

Whether to shuffle data.

---

**Returns:** `Tuple[DataLoader, DistributedSampler | None]`

Tuple of (DataLoader, DistributedSampler or None).

```python
nemo_automodel.components.datasets.diffusion.mock_dataloader.mock_collate_fn(
    batch
)
```

Collate function for mock dataset, matching the real collate\_fn behavior.