Repository: Lightning-AI/pytorch-lightning
File: tests/core/test_datamodules.py (Python)

Code before:
    def setup(self, stage: Optional[str] = None):

        mnist_full = TrialMNIST(
            root=self.data_dir, train=True, num_samples=64, download=True
        )
        self.mnist_train, self.mnist_val = random_split(mnist_full, [128, 64])
        self.dims = self.mnist_train[0][0].shape

    def train_dataloader(self):
        assert self.trainer.current_epoch not in self._epochs_called_for
        self._epochs_called_for.append(self.trainer.current_epoch)

        return DataLoader(self.mnist_train, batch_size=4)


def test_dm_reload_dataloaders_every_epoch(tmpdir):
    """Test datamodule, where trainer argument
    reload_dataloaders_every_epoch is set to True/False"""

    dm = CustomMNISTDataModule(tmpdir)

    model = EvalModelTemplate()
    model.validation_step = None
    model.validation_step_end = None
    model.validation_epoch_end = None
    model.test_step = None
    model.test_step_end = None
    model.test_epoch_end = None

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=2,
        limit_train_batches=0.01,
        reload_dataloaders_every_epoch=True,
    )
    trainer.fit(model, dm)


Review comment (performance):
we can simplify this:

```suggestion
    assert all(batch_gpu.samples == 1)
    assert all(batch_gpu.targets == 2)
```

assuming you don't need to test the shape ~or dtype~

Code after:
            batch.targets = batch.targets.to(device)
            return batch

    dm = CurrentTestDM()
    model = BoringModel()

    batch = CustomBatch((torch.zeros(5, 32), torch.ones(5, 1, dtype=torch.long)))

    trainer = Trainer(gpus=1)
    # running .fit() would require us to implement custom data loaders, we mock the model reference instead
    get_module_mock.return_value = model
    if is_overridden('transfer_batch_to_device', dm):
        model.transfer_batch_to_device = dm.transfer_batch_to_device

    model.on_before_batch_transfer = dm.on_before_batch_transfer
    model.transfer_batch_to_device = dm.transfer_batch_to_device
    model.on_after_batch_transfer = dm.on_after_batch_transfer

    batch_gpu = trainer.accelerator_backend.batch_to_device(batch, expected_device)

    assert dm.on_before_batch_transfer_hook_rank == 0
    assert dm.transfer_batch_to_device_hook_rank == 1
    assert dm.on_after_batch_transfer_hook_rank == 2
    assert batch_gpu.samples.device == batch_gpu.targets.device == expected_device
    assert torch.allclose(batch_gpu.samples.cpu(), torch.ones(5, 32))
    assert torch.allclose(batch_gpu.targets.cpu(), torch.ones(5, 1, dtype=torch.long) * 2)


def test_dm_reload_dataloaders_every_epoch(tmpdir):
    """Test datamodule, where trainer argument
    reload_dataloaders_every_epoch is set to True/False"""

    class CustomBoringDataModule(BoringDataModule):

        def __init__(self):
            super().__init__()
            self._epochs_called_for = []

        def train_dataloader(self):
            assert self.trainer.current_epoch not in self._epochs_called_for
            self._epochs_called_for.append(self.trainer.current_epoch)
            return super().train_dataloader()

    dm = CustomBoringDataModule()
    model = BoringModel()

    model.validation_step = None
    model.validation_step_end = None
    model.validation_epoch_end = None
    model.test_step = None
    model.test_step_end = None