Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/core/hooks.py (Python)

Code before:
            A reference to the data on the new device.

        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).

        Note:
            This hook only runs on single GPU training (no data-parallel). If you need multi-GPU support
            for your custom batch objects, you need to define your custom
            :class:`~torch.nn.parallel.DistributedDataParallel` or
            :class:`~pytorch_lightning.overrides.data_parallel.LightningDistributedDataParallel` and
            override :meth:`~pytorch_lightning.core.lightning.LightningModule.configure_ddp`.

        See Also:
            - :func:`~pytorch_lightning.utilities.apply_func.move_data_to_device`
            - :func:`~pytorch_lightning.utilities.apply_func.apply_to_collection`
        """
        return move_data_to_device(batch, device)

    def before_batch_transfer(self, batch):
        return batch

    def after_batch_transfer(self, batch):
        return batch

    def prepare_batch_for_transfer(self, batch, device):
        batch = self.before_batch_transfer(batch)
        batch = self.transfer_batch_to_device(batch, device)
        batch = self.after_batch_transfer(batch)
        return batch


Review comment (performance):
this name is not good enough for what the code below describes. 
Why are these hooks needed if they immediately follow each other. 
Right now, you could do the same by overriding only transfer_batch_to_device and do the before- and after- code in there. 

Code after:

        Returns:
            A reference to the data on the new device.

        Example::

            def transfer_batch_to_device(self, batch, device):
                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)
                else:
                    batch = super().transfer_batch_to_device(data, device)
                return batch

        See Also:
            - :meth:`move_data_to_device`
            - :meth:`apply_to_collection`
        """
        device = device or self.device
        return move_data_to_device(batch, device)

    def on_before_batch_transfer(self, batch):
        """
        Override to alter or apply batch augmentations to your batch before it is transferred to the device.

        Note:
            This hook only runs on single GPU training and DDP (no data-parallel).

        Args:
            batch: A batch of data that needs to be altered or augmented.

        Returns:
            A batch of data

        Example::

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

        See Also:
            - :meth:`on_after_batch_transfer`
            - :meth:`transfer_batch_to_device`
        """
        return batch

    def on_after_batch_transfer(self, batch):
        """
        Override to alter or apply batch augmentations to your batch after it is transferred to the device.