---
title: DataHooks
description: Hooks to be used for data related stuff.
canonical: >-
  https://lightning.ai/docs/pytorch/stable/optional-api/api/lightning.pytorch.core.hooks.DataHooks
---
# DataHooks {#datahooks}

### `class lightning.pytorch.core.hooks.DataHooks` {#lightning.pytorch.core.hooks.DataHooks}

Bases: [`object`](https://docs.python.org/3/builtins/functions.html#object "(in Python v3.15)")

Hooks to be used for data related stuff.

#### `prepare_data_per_node` {#lightning.pytorch.core.hooks.DataHooks.prepare_data_per_node}

If True, each LOCAL\_RANK=0 will call prepare data. Otherwise only NODE\_RANK=0, LOCAL\_RANK=0 will prepare data.

#### `allow_zero_length_dataloader_with_multiple_devices` {#lightning.pytorch.core.hooks.DataHooks.allow_zero_length_dataloader_with_multiple_devices}

If True, dataloader with zero length within local rank is allowed. Default value is False.

#### `on_after_batch_transfer(batch, dataloader_idx)` {#lightning.pytorch.core.hooks.DataHooks.on_after_batch_transfer}

Override to alter or apply batch augmentations to your batch after it is transferred to the device.

Note

To check the current state of execution of this hook you can use `self.trainer.training/testing/validating/predicting` so that you can add different logic as per your requirement.

Parameters:

*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – A batch of data that needs to be altered or augmented.
    
*   **dataloader\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – The index of the dataloader to which the batch belongs.
    

Return type:

[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")

Returns:

A batch of data

Example:

```python
def on_after_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = gpu_transforms(batch['x'])
    return batch
```

See also

*   [`on_before_batch_transfer()`](#lightning.pytorch.core.hooks.DataHooks.on_before_batch_transfer "lightning.pytorch.core.hooks.DataHooks.on_before_batch_transfer")
    
*   [`transfer_batch_to_device()`](#lightning.pytorch.core.hooks.DataHooks.transfer_batch_to_device "lightning.pytorch.core.hooks.DataHooks.transfer_batch_to_device")
    

#### `on_before_batch_transfer(batch, dataloader_idx)` {#lightning.pytorch.core.hooks.DataHooks.on_before_batch_transfer}

Override to alter or apply batch augmentations to your batch before it is transferred to the device.

Note

To check the current state of execution of this hook you can use `self.trainer.training/testing/validating/predicting` so that you can add different logic as per your requirement.

Parameters:

*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – A batch of data that needs to be altered or augmented.
    
*   **dataloader\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – The index of the dataloader to which the batch belongs.
    

Return type:

[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")

Returns:

A batch of data

Example:

```python
def on_before_batch_transfer(self, batch, dataloader_idx):
    batch['x'] = transforms(batch['x'])
    return batch
```

See also

*   [`on_after_batch_transfer()`](#lightning.pytorch.core.hooks.DataHooks.on_after_batch_transfer "lightning.pytorch.core.hooks.DataHooks.on_after_batch_transfer")
    
*   [`transfer_batch_to_device()`](#lightning.pytorch.core.hooks.DataHooks.transfer_batch_to_device "lightning.pytorch.core.hooks.DataHooks.transfer_batch_to_device")
    

#### `predict_dataloader()` {#lightning.pytorch.core.hooks.DataHooks.predict_dataloader}

An iterable or collection of iterables specifying prediction samples.

For more information about multiple dataloaders, see this [section](/docs/pytorch/2.6.6/reference/data/iterables.md#multiple-dataloaders).

It’s recommended that all data downloads and preparation happen in [`prepare_data()`](#lightning.pytorch.core.hooks.DataHooks.prepare_data "lightning.pytorch.core.hooks.DataHooks.prepare_data").

*   [`predict()`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer.predict "lightning.pytorch.trainer.trainer.Trainer.predict")
    
*   [`prepare_data()`](#lightning.pytorch.core.hooks.DataHooks.prepare_data "lightning.pytorch.core.hooks.DataHooks.prepare_data")
    
*   [`setup()`](#lightning.pytorch.core.hooks.DataHooks.setup "lightning.pytorch.core.hooks.DataHooks.setup")
    

Note

Lightning tries to add the correct sampler for distributed and arbitrary hardware There is no need to set it yourself.

Return type:

[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")

Returns:

A [`torch.utils.data.DataLoader`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "(in PyTorch v2.15)") or a sequence of them specifying prediction samples.

#### `prepare_data()` {#lightning.pytorch.core.hooks.DataHooks.prepare_data}

Use this to download and prepare data. Downloading and saving data with multiple processes (distributed settings) will result in corrupted data. Lightning ensures this method is called only within a single process, so you can safely add your downloading logic within. :rtype: [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

Warning

DO NOT set state to the model (use `setup` instead) since this is NOT called on every device

Example:

```python
def prepare_data(self):
    # good
    download_data()
    tokenize()
    etc()

    # bad
    self.split = data_split
    self.some_state = some_other_state()
```

In a distributed environment, `prepare_data` can be called in two ways (using [prepare\_data\_per\_node](/docs/pytorch/2.6.6/core-api/lightning_module.md#prepare-data-per-node))

1.  Once per node. This is the default and is only called on LOCAL\_RANK=0.
    
2.  Once in total. Only called on GLOBAL\_RANK=0.
    

Example:

```python
# DEFAULT
# called once per node on LOCAL_RANK=0 of that node
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = True


# call on GLOBAL_RANK=0 (great for shared file systems)
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = False
```

This is called before requesting the dataloaders:

```python
model.prepare_data()
initialize_distributed()
model.setup(stage)
model.train_dataloader()
model.val_dataloader()
model.test_dataloader()
model.predict_dataloader()
```

#### `setup(stage)` {#lightning.pytorch.core.hooks.DataHooks.setup}

Called at the beginning of fit (train + validate), validate, test, or predict. This is a good hook when you need to build models dynamically or adjust something about them. This hook is called on every process when using DDP.

Parameters:

**stage** ([`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) – either `'fit'`, `'validate'`, `'test'`, or `'predict'`

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

Example:

```python
class LitModel(...):
    def __init__(self):
        self.l1 = None

    def prepare_data(self):
        download_data()
        tokenize()

        # don't do this
        self.something = else

    def setup(self, stage):
        data = load_data(...)
        self.l1 = nn.Linear(28, data.num_classes)
```

#### `teardown(stage)` {#lightning.pytorch.core.hooks.DataHooks.teardown}

Called at the end of fit (train + validate), validate, test, or predict.

Parameters:

**stage** ([`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")) – either `'fit'`, `'validate'`, `'test'`, or `'predict'`

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `test_dataloader()` {#lightning.pytorch.core.hooks.DataHooks.test_dataloader}

An iterable or collection of iterables specifying test samples.

For more information about multiple dataloaders, see this [section](/docs/pytorch/2.6.6/reference/data/iterables.md#multiple-dataloaders).

For data processing use the following pattern: :rtype: [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")

*   download in [`prepare_data()`](#lightning.pytorch.core.hooks.DataHooks.prepare_data "lightning.pytorch.core.hooks.DataHooks.prepare_data")
    
*   process and split in [`setup()`](#lightning.pytorch.core.hooks.DataHooks.setup "lightning.pytorch.core.hooks.DataHooks.setup")
    

However, the above are only necessary for distributed processing.

Warning

do not assign state in prepare\_data

*   [`test()`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer.test "lightning.pytorch.trainer.trainer.Trainer.test")
    
*   [`prepare_data()`](#lightning.pytorch.core.hooks.DataHooks.prepare_data "lightning.pytorch.core.hooks.DataHooks.prepare_data")
    
*   [`setup()`](#lightning.pytorch.core.hooks.DataHooks.setup "lightning.pytorch.core.hooks.DataHooks.setup")
    

Note

Lightning tries to add the correct sampler for distributed and arbitrary hardware. There is no need to set it yourself.

Note

If you don’t need a test dataset and a `test_step()`, you don’t need to implement this method.

#### `train_dataloader()` {#lightning.pytorch.core.hooks.DataHooks.train_dataloader}

An iterable or collection of iterables specifying training samples.

For more information about multiple dataloaders, see this [section](/docs/pytorch/2.6.6/reference/data/iterables.md#multiple-dataloaders).

The dataloader you return will not be reloaded unless you set [`reload_dataloaders_every_n_epochs`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer.params.reload_dataloaders_every_n_epochs "lightning.pytorch.trainer.trainer.Trainer") to a positive integer.

For data processing use the following pattern: :rtype: [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")

*   download in [`prepare_data()`](#lightning.pytorch.core.hooks.DataHooks.prepare_data "lightning.pytorch.core.hooks.DataHooks.prepare_data")
    
*   process and split in [`setup()`](#lightning.pytorch.core.hooks.DataHooks.setup "lightning.pytorch.core.hooks.DataHooks.setup")
    

However, the above are only necessary for distributed processing.

Warning

do not assign state in prepare\_data

*   [`fit()`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer.fit "lightning.pytorch.trainer.trainer.Trainer.fit")
    
*   [`prepare_data()`](#lightning.pytorch.core.hooks.DataHooks.prepare_data "lightning.pytorch.core.hooks.DataHooks.prepare_data")
    
*   [`setup()`](#lightning.pytorch.core.hooks.DataHooks.setup "lightning.pytorch.core.hooks.DataHooks.setup")
    

Note

Lightning tries to add the correct sampler for distributed and arbitrary hardware. There is no need to set it yourself.

#### `transfer_batch_to_device(batch, device, dataloader_idx)` {#lightning.pytorch.core.hooks.DataHooks.transfer_batch_to_device}

Override this hook if your [`DataLoader`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "(in PyTorch v2.15)") returns tensors wrapped in a custom data structure.

The data types listed below (and any arbitrary nesting of them) are supported out of the box:

*   [`torch.Tensor`](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "(in PyTorch v2.15)") or anything that implements .to(…)
    
*   [`list`](https://docs.python.org/3/builtins/stdtypes.html#list "(in Python v3.15)")
    
*   [`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")
    
*   [`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple "(in Python v3.15)")
    

For anything else, you need to define how the data is moved to the target device (CPU, GPU, TPU, …).

Note

This hook should only transfer the data and not modify it, nor should it move the data to any other device than the one passed in as argument (unless you know what you are doing). To check the current state of execution of this hook you can use `self.trainer.training/testing/validating/predicting` so that you can add different logic as per your requirement.

Parameters:

*   **batch** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – A batch of data that needs to be transferred to a new device.
    
*   **device** ([`device`](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "(in PyTorch v2.15)")) – The target device as defined in PyTorch.
    
*   **dataloader\_idx** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – The index of the dataloader to which the batch belongs.
    

Return type:

[`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")

Returns:

A reference to the data on the new device.

Example:

```python
def transfer_batch_to_device(self, batch, device, dataloader_idx):
    if isinstance(batch, CustomBatch):
        # move all tensors in your custom data structure to the device
        batch.samples = batch.samples.to(device)
        batch.targets = batch.targets.to(device)
    elif dataloader_idx == 0:
        # skip device transfer for the first dataloader or anything you wish
        pass
    else:
        batch = super().transfer_batch_to_device(batch, device, dataloader_idx)
    return batch
```

See also

*   `move_data_to_device()`
    
*   `apply_to_collection()`
    

#### `val_dataloader()` {#lightning.pytorch.core.hooks.DataHooks.val_dataloader}

An iterable or collection of iterables specifying validation samples.

For more information about multiple dataloaders, see this [section](/docs/pytorch/2.6.6/reference/data/iterables.md#multiple-dataloaders).

The dataloader you return will not be reloaded unless you set [`reload_dataloaders_every_n_epochs`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer.params.reload_dataloaders_every_n_epochs "lightning.pytorch.trainer.trainer.Trainer") to a positive integer.

It’s recommended that all data downloads and preparation happen in [`prepare_data()`](#lightning.pytorch.core.hooks.DataHooks.prepare_data "lightning.pytorch.core.hooks.DataHooks.prepare_data"). :rtype: [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")

*   [`fit()`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer.fit "lightning.pytorch.trainer.trainer.Trainer.fit")
    
*   [`validate()`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer.validate "lightning.pytorch.trainer.trainer.Trainer.validate")
    
*   [`prepare_data()`](#lightning.pytorch.core.hooks.DataHooks.prepare_data "lightning.pytorch.core.hooks.DataHooks.prepare_data")
    
*   [`setup()`](#lightning.pytorch.core.hooks.DataHooks.setup "lightning.pytorch.core.hooks.DataHooks.setup")
    

Note

Lightning tries to add the correct sampler for distributed and arbitrary hardware There is no need to set it yourself.

Note

If you don’t need a validation dataset and a `validation_step()`, you don’t need to implement this method.

---

**Navigation**

- Parent: [Api](/docs/pytorch/2.6.6/optional-api/api.md)