diff --git a/docs/Classification.ipynb b/docs/Classification.ipynb
index 06f02b6fefcb4839261f194520ec1b2131aba1a6..0606156c72ebe4390b27ac32b8d245e8acac61bd 100644
--- a/docs/Classification.ipynb
+++ b/docs/Classification.ipynb
@@ -16,9 +16,9 @@
     "import pandas as pd\n",
     "import matplotlib.pyplot as plt\n",
     "import seaborn as sns\n",
-    "from keras.models import Sequential\n",
-    "from keras.layers import Dense, Dropout\n",
-    "from keras.optimizers import SGD\n",
+    "from tensorflow.keras.models import Sequential\n",
+    "from tensorflow.keras.layers import Dense, Dropout\n",
+    "from tensorflow.keras.optimizers import SGD\n",
     "\n",
     "from freeforestml import Variable, Process, Cut, \\\n",
     "                    HepNet, ClassicalCV, EstimatorNormalizer, \\\n",
diff --git a/freeforestml/helpers.py b/freeforestml/helpers.py
index ba96437445cb567159ac9912150919d344770db2..3ba1da5a3d6c5717efe41d2d00a62b8793db10e9 100644
--- a/freeforestml/helpers.py
+++ b/freeforestml/helpers.py
@@ -1,12 +1,12 @@
 import base64
-import cloudpickle
+import dill
 
 
 def python_to_str(obj):
     """
     Convert an arbitrary python object into a string and encode it in base64.
     """
-    obj_string = cloudpickle.dumps(obj)
+    obj_string = dill.dumps(obj)
     obj_base64 = base64.b64encode(obj_string).decode()
     return obj_base64
 
@@ -16,5 +16,5 @@ def str_to_python(string):
     Reverse of the python_to_str() function.
     """
     obj_string = base64.b64decode(string)
-    obj = cloudpickle.loads(obj_string)
+    obj = dill.loads(obj_string)
     return obj
diff --git a/freeforestml/model.py b/freeforestml/model.py
index e1da2df4d61820c994fd2f23141e7bb40219f93c..de45ed5bda1113c178f76554caca077e7554f427 100644
--- a/freeforestml/model.py
+++ b/freeforestml/model.py
@@ -7,7 +7,7 @@ import json
 
 import numpy as np
 import pandas as pd
-import keras
+import tensorflow
 
 from freeforestml.variable import Variable
 from freeforestml.helpers import python_to_str, str_to_python
@@ -796,7 +796,7 @@ class HepNet:
                 else:
                     path_token.insert(-1, f"fold_{fold_i}")
 
-                model = keras.models.load_model(".".join(path_token))
+                model = tensorflow.keras.models.load_model(".".join(path_token))
                 instance.models.append(model)
 
         # load normalizer
diff --git a/freeforestml/tests/test_model.py b/freeforestml/tests/test_model.py
index 3841b60cc38e70c3ac37d0fb68217d6e0aeb5b28..d7c9b809bd5fd749bedd6d0a1c92e24f172ced70 100644
--- a/freeforestml/tests/test_model.py
+++ b/freeforestml/tests/test_model.py
@@ -5,9 +5,9 @@ import unittest
 import math
 
 import pandas as pd
-from keras.models import Sequential
-from keras.layers import Dense, Dropout
-from keras.optimizers import SGD
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Dense, Dropout
+from tensorflow.keras.optimizers import SGD
 
 from freeforestml.model import CrossValidator, ClassicalCV, MixedCV, \
                           Normalizer, EstimatorNormalizer, \
diff --git a/freeforestml/variable.py b/freeforestml/variable.py
index bbdbb98fb503ef6f4c29100c45ee62b21b8a0b7d..35857c9037991c40bc605dc5fda93ddcd9d78b69 100644
--- a/freeforestml/variable.py
+++ b/freeforestml/variable.py
@@ -2,7 +2,7 @@ from abc import ABC, abstractmethod
 
 import base64
 import os
-import cloudpickle
+import dill
 import numpy as np
 import h5py
 
diff --git a/requirements.txt b/requirements.txt
index 2ca7dd81af6eac3dfa4587b7ea9d7300d86ca091..9ba6bb867a8d52a95b25cf3f0a00f499771b5135 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
-cloudpickle
+dill
 h5py
 numpy
 matplotlib
@@ -8,7 +8,6 @@ pandas
 pylorentz
 atlasify
 lxml
-keras
 tensorflow
 dask[complete]
 uhepp
diff --git a/setup.py b/setup.py
index 984ca1daa04ecfe307266c57cb133b4985a102a7..9e00f2068d8db5533d6a4c83d59517732aed982d 100644
--- a/setup.py
+++ b/setup.py
@@ -12,7 +12,7 @@ from setuptools import setup
 setup(name='freeforestml',
       version='0.0.0',  # Also change in module
       packages=["freeforestml", "freeforestml.tests"],
-      install_requires=["cloudpickle",
+      install_requires=["dill",
                         "h5py",
                         "numpy",
                         "matplotlib",
@@ -22,7 +22,6 @@ setup(name='freeforestml',
                         "pylorentz",
                         "atlasify>=0.2.0",
                         "tensorflow",
-                        "keras",
                         "dask",
                         "dask[complete]",
                         "uhepp",