Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/trainer/connectors/data_connector.py (Python)

Code before:

        if test_dataloaders is not None:
            model.test_dataloader = _PatchDataLoader(test_dataloaders)

    def attach_datamodule(self, model, datamodule, stage):

        # We use datamodule if it's been provided on .fit or .test, otherwise we check model for it
        datamodule = datamodule or getattr(model, 'datamodule', None)

        # If we have a datamodule, attach necessary hooks + dataloaders
        if datamodule:

            # Override loader hooks
            if is_overridden('train_dataloader', datamodule):
                model.train_dataloader = datamodule.train_dataloader
            if is_overridden('val_dataloader', datamodule):
                model.val_dataloader = datamodule.val_dataloader
            if is_overridden('test_dataloader', datamodule):
                model.test_dataloader = datamodule.test_dataloader

            # TODO check this before merge
            # Override prepare_batch_for_transfer if dataset-specific to_device logic has been defined in datamodule
            model.prepare_batch_for_transfer = datamodule.prepare_batch_for_transfer

            # TODO remove this after all the changes are done
            model.transfer_batch_to_device = datamodule.transfer_batch_to_device

            self.trainer.datamodule = datamodule


class _PatchDataLoader(object):
    r"""
    Callable object for patching dataloaders passed into trainer.fit().
    Use this class to override model.*_dataloader() and be pickle-compatible.

    Args:
        dataloader: Dataloader object to return when called.

    """

    def __init__(self, dataloader: Union[List[DataLoader], DataLoader]):
        self.dataloader = dataloader

        # cannot pickle __code__ so cannot verify if PatchDataloader
        # exists which shows dataloader methods have been overwritten.
        # so, we hack it by using the string representation
        self.patch_loader_code = str(self.__call__.__code__)

    def __call__(self) -> Union[List[DataLoader], DataLoader]:
        return self.dataloader


Review comment (suggestion):
I understand what you try to achieve here. It is done in other places. 

However, I am not a fan of patching the model methods through the Trainer. 
This is already done with other methods, like the dataloaders :(
The bad thing I see here is that if I have a method `module.method`,
during `trainer.fit(model)` Trainer replaces `module.method = method2` without me knowing, 
now I expect to be able to run `model.method` after fit but it behaves differently. Why? because my model DEFINITION was altered behind the scenes.
In my view, this is currently an issue in Lightning, and I will bring it up again after v1.0. I would very much prefer not to do this patching. After all, this LightningDataModule is supposed to be separate from LightningModule, that's why it was made. The trainer should USE the methods in it for training, but NOT substitute methods in my LightningModule. I hope you can follow my thoughts. 

Code after:
    ):
        # when dataloader is passed via fit, patch the train_dataloader
        # functions to overwrite with these implementations
        if train_dataloader is not None:
            model.train_dataloader = _PatchDataLoader(train_dataloader)

        if val_dataloaders is not None:
            model.val_dataloader = _PatchDataLoader(val_dataloaders)

        if test_dataloaders is not None:
            model.test_dataloader = _PatchDataLoader(test_dataloaders)

        if predict_dataloaders is not None:
            model.predict_dataloader = _PatchDataLoader(predict_dataloaders)

    def attach_datamodule(self, model, datamodule: Optional[LightningDataModule], stage: str) -> None:
        # Todo: required argument `stage` is not used

        # We use datamodule if it's been provided on .fit or .test, otherwise we check model for it
        datamodule = datamodule or getattr(model, 'datamodule', None)

        # If we have a datamodule, attach necessary hooks + dataloaders
        if datamodule:

            # Override loader hooks
            if is_overridden('train_dataloader', datamodule):
                model.train_dataloader = datamodule.train_dataloader
            if is_overridden('val_dataloader', datamodule):
                model.val_dataloader = datamodule.val_dataloader
            if is_overridden('test_dataloader', datamodule):
                model.test_dataloader = datamodule.test_dataloader
            if is_overridden('predict_dataloader', datamodule):
                model.predict_dataloader = datamodule.predict_dataloader

            # Override data transfer hooks if dataset-specific to_device logic has been defined in datamodule
            if is_overridden('on_before_batch_transfer', datamodule):
                model.on_before_batch_transfer = datamodule.on_before_batch_transfer
            if is_overridden('transfer_batch_to_device', datamodule):
                model.transfer_batch_to_device = datamodule.transfer_batch_to_device
            if is_overridden('on_after_batch_transfer', datamodule):
                model.on_after_batch_transfer = datamodule.on_after_batch_transfer

            self.trainer.datamodule = datamodule
            datamodule.trainer = self.trainer


class _PatchDataLoader(object):
    r"""
    Callable object for patching dataloaders passed into trainer.fit().
    Use this class to override model.*_dataloader() and be pickle-compatible.