Replace __reduce__ with __getstate__

Author: carmocca

pytorch_lightning/core/datamodule.py

@@ -326,11 +326,8 @@ def test_dataloader():
             datamodule.test_dataloader = test_dataloader
         return datamodule
 
-    def __new__(cls, *args: Any, **kwargs: Any) -> 'LightningDataModule':
+    def __new__(cls, *_: Any, **__: Any) -> 'LightningDataModule':
         obj = super().__new__(cls)
-        # save `args` and `kwargs` for `__reduce__`
-        obj.__args = args
-        obj.__kwargs = kwargs
         # track `DataHooks` calls and run `prepare_data` only on rank zero
         obj.prepare_data = cls._track_data_hook_calls(obj, rank_zero_only(obj.prepare_data))
         obj.setup = cls._track_data_hook_calls(obj, obj.setup)
@@ -381,6 +378,9 @@ def wrapped_fn(*args: str, **kwargs: Optional[str]) -> Any:
 
         return wrapped_fn
 
-    def __reduce__(self) -> Tuple[type, tuple, dict]:
+    def __getstate__(self) -> dict:
         # avoids _pickle.PicklingError: Can't pickle <...>: it's not the same object as <...>
-        return self.__class__, self.__args, self.__kwargs
+        d = self.__dict__.copy()
+        for fn in ("prepare_data", "setup", "teardown"):
+            del d[fn]
+        return d