replace all pickle calls with dill calls

mlair/data_handler/default_data_handler.py

@@ -8,6 +8,7 @@ import gc
 import logging
 import os
 import pickle
+import dill
 import shutil
 from functools import reduce
 from typing import Tuple, Union, List
@@ -86,7 +87,7 @@ class DefaultDataHandler(AbstractDataHandler):
             data = {"X": self._X, "Y": self._Y, "X_extreme": self._X_extreme, "Y_extreme": self._Y_extreme}
             data = self._force_dask_computation(data)
             with open(self._save_file, "wb") as f:
-                pickle.dump(data, f)
+                dill.dump(data, f)
             logging.debug(f"save pickle data to {self._save_file}")
             self._reset_data()
 
@@ -101,7 +102,7 @@ class DefaultDataHandler(AbstractDataHandler):
     def _load(self):
         try:
             with open(self._save_file, "rb") as f:
-                data = pickle.load(f)
+                data = dill.load(f)
             logging.debug(f"load pickle data from {self._save_file}")
             self._X, self._Y = data["X"], data["Y"]
             self._X_extreme, self._Y_extreme = data["X_extreme"], data["Y_extreme"]

mlair/data_handler/iterator.py

@@ -9,6 +9,7 @@ import math
 import os
 import shutil
 import pickle
+import dill
 from typing import Tuple, List
 
 
@@ -109,7 +110,7 @@ class KerasIterator(keras.utils.Sequence):
         """Load pickle data from disk."""
         file = self._path % index
         with open(file, "rb") as f:
-            data = pickle.load(f)
+            data = dill.load(f)
         return data["X"], data["Y"]
 
     @staticmethod
@@ -167,7 +168,7 @@ class KerasIterator(keras.utils.Sequence):
         data = {"X": X, "Y": Y}
         file = self._path % index
         with open(file, "wb") as f:
-            pickle.dump(data, f)
+            dill.dump(data, f)
 
     def _get_number_of_mini_batches(self, number_of_samples: int) -> int:
         """Return number of mini batches as the floored ration of number of samples to batch size."""