> 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                                                               |
| ------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------- |
| [`MockWanDataset`](#nemo_automodel-components-datasets-diffusion-mock_dataloader-MockWanDataset) | Mock dataset that generates random data matching WAN 2.1 expected format. |

### 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.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:**

Number of samples in the dataset.

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

Number of temporal latent frames.

Height of spatial latents.

Width of spatial latents.

Length of text sequence.

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

Device to place tensors on.

```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
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, typing.Optional[torch.utils.data.DistributedSampler]]
```

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:**

Data parallel rank.

Data parallel world size.

Batch size per GPU.

Number of dataloader workers.

Device to place tensors on.

Number of samples in mock dataset.

Number of latent channels (default: 16).

Number of temporal latent frames.

Height of spatial latents.

Width of spatial latents.

Length of text sequence.

Dimension of text embeddings.

Whether to shuffle data.

**Returns:** `Tuple[DataLoader, Optional[DistributedSampler]]`

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.