Fix PL compatibility

CHANGELOG.md

@@ -78,6 +78,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed a bug where `drop_last` would be set to True during prediction and testing ([#671](https://github.com/PyTorchLightning/lightning-flash/pull/671))
 
+- Fixed a bug where flash was not compatible with pytorch-lightning >= 1.4.3 ([#690](https://github.com/PyTorchLightning/lightning-flash/pull/690))
+
 ## [0.4.0] - 2021-06-22
 
 ### Added

flash/core/data/data_pipeline.py

@@ -14,7 +14,7 @@
 import functools
 import inspect
 import weakref
-from typing import Any, Callable, Dict, Optional, Sequence, Set, Tuple, Type, TYPE_CHECKING
+from typing import Any, Callable, Dict, Optional, Sequence, Set, Tuple, Type, TYPE_CHECKING, Union
 
 import torch
 from pytorch_lightning.trainer.connectors.data_connector import _PatchDataLoader
@@ -29,6 +29,7 @@
 from flash.core.data.process import DefaultPreprocess, Deserializer, Postprocess, Preprocess, Serializer
 from flash.core.data.properties import ProcessState
 from flash.core.data.utils import _POSTPROCESS_FUNCS, _PREPROCESS_FUNCS, _STAGES_PREFIX
+from flash.core.utilities.imports import _PL_GREATER_EQUAL_1_4_3
 
 if TYPE_CHECKING:
     from flash.core.model import Task
@@ -331,6 +332,16 @@ def _get_dataloader(model: "Task", loader_name: str) -> Tuple[DataLoader, str]:
 
         return dataloader, attr_name
 
+    @staticmethod
+    def _patch_dataloader(model: "Task", dataloader: Union[Callable, DataLoader], stage: RunningStage):
+        if isinstance(dataloader, DataLoader):
+            if _PL_GREATER_EQUAL_1_4_3:
+                dataloader = _PatchDataLoader(dataloader, _STAGES_PREFIX[stage])
+                dataloader.patch(model)
+            else:
+                dataloader = _PatchDataLoader(dataloader)
+        return dataloader
+
     @staticmethod
     def _set_loader(model: "Task", loader_name: str, new_loader: DataLoader) -> None:
         """This function is used to set the loader to model and/or datamodule."""
@@ -405,8 +416,7 @@ def _attach_preprocess_to_model(
                 if not was_seq:
                     dataloader = dataloader[0]
 
-                if isinstance(dataloader, DataLoader):
-                    dataloader = _PatchDataLoader(dataloader)
+                dataloader = self._patch_dataloader(model, dataloader, stage)
 
                 self._set_loader(model, whole_attr_name, dataloader)
 
@@ -535,8 +545,7 @@ def _detach_preprocessing_from_model(self, model: "Task", stage: Optional[Runnin
             if not was_seq:
                 dataloader = dataloader[0]
 
-            if isinstance(dataloader, DataLoader):
-                dataloader = _PatchDataLoader(dataloader)
+            dataloader = self._patch_dataloader(model, dataloader, stage)
 
             self._set_loader(model, whole_attr_name, dataloader)
 

flash/core/utilities/imports.py

@@ -121,6 +121,7 @@ class Image(metaclass=MetaImage):
 
 if Version:
     _TORCHVISION_GREATER_EQUAL_0_9 = _compare_version("torchvision", operator.ge, "0.9.0")
+    _PL_GREATER_EQUAL_1_4_3 = _compare_version("pytorch_lightning", operator.ge, "1.4.3")
 
 _TEXT_AVAILABLE = all(
     [