Fix _PatchDataLoader

Lightning-AI · Aug 22, 2021 · b947f07 · b947f07
1 parent b13ca7d
commit b947f07
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/pytorch_lightning/trainer/connectors/data_connector.py b/pytorch_lightning/trainer/connectors/data_connector.py
@@ -190,10 +190,14 @@ def __init__(self, dataloader: Union[TRAIN_DATALOADERS, EVAL_DATALOADERS], stage
     def __call__(self) -> Union[TRAIN_DATALOADERS, EVAL_DATALOADERS]:
         return self.dataloader
 
+    @property
+    def __name__(self) -> str:
+        return self.stage + "_dataloader"
+
     def patch(self, model: "pl.LightningModule") -> None:
-        self._old_loader = getattr(model, self.stage + "_dataloader")
-        setattr(model, self.stage + "_dataloader", self)
+        self._old_loader = getattr(model, self.__name__)
+        setattr(model, self.__name__, self)
 
     def unpatch(self, model: "pl.LightningModule") -> None:
-        setattr(model, self.stage + "_dataloader", self._old_loader)
+        setattr(model, self.__name__, self._old_loader)
         self._old_loader = None
diff --git a/tests/trainer/test_dataloaders.py b/tests/trainer/test_dataloaders.py
@@ -1477,6 +1477,10 @@ def __init__(self, loader):
         def __call__(self):
             return self.loader
 
+        def __name__(self):
+            # `_PatchDataLoader` requires this when passed to `trainer.call_hook`
+            return "foo"
+
     class TestModel(BoringModel):
         def __init__(self):
             super().__init__()