diff --git a/.github/workflows/pythonapp.yml b/.github/workflows/pythonapp.yml
index 9aba762..99eb7b6 100644
--- a/.github/workflows/pythonapp.yml
+++ b/.github/workflows/pythonapp.yml
@@ -7,9 +7,9 @@ jobs:
 
     runs-on: ubuntu-latest
     strategy:
-        max-parallel: 2
+        max-parallel: 3
         matrix:
-            python-version: [3.6, 3.7]
+            python-version: [3.6, 3.7, 3.8]
 
     steps:
     - uses: actions/checkout@v1
@@ -20,6 +20,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
+        pip install setuptools --upgrade
         pip install .
     - name: Lint with pylint
       run: |
@@ -31,5 +32,4 @@ jobs:
     - name: Test with pytest
       run: |
         pip install pytest
-        pip install tensorflow # Keras does not install any backend by default?
         pytest
diff --git a/.readthedocs.yml b/.readthedocs.yml
index 3107642..536e072 100644
--- a/.readthedocs.yml
+++ b/.readthedocs.yml
@@ -6,4 +6,9 @@ sphinx:
 
 python:
   version: 3.7
-  pip_install: True
+  install:
+      - method: pip
+        path: .
+        extra_requirements:
+            - docs
+  system_packages: true
diff --git a/doc/source/conf.py b/doc/source/conf.py
index da0e055..b75ad9a 100644
--- a/doc/source/conf.py
+++ b/doc/source/conf.py
@@ -26,6 +26,7 @@
 
 # The full version, including alpha/beta/rc tags
 release = evolutionary_keras .__version__
+autodoc_mock_imports = ['tensorflow']
 
 
 # -- General configuration ---------------------------------------------------
diff --git a/examples/test.py b/examples/test.py
index 662d6c2..834ed02 100644
--- a/examples/test.py
+++ b/examples/test.py
@@ -1,21 +1,25 @@
 """ Example implementation of evolutionary_keras
 """
 import sys
-import keras
-from keras.datasets import mnist
-from keras import backend as K
-from keras.layers import Dense, Input, Flatten
-
+import tensorflow as tf
+from tensorflow.keras.datasets import mnist
+from tensorflow.keras import backend as K
+from tensorflow.keras.layers import Dense, Input, Flatten
 from evolutionary_keras.models import EvolModel
 import evolutionary_keras.optimizers
 
+import logging
+logging.basicConfig(level = logging.INFO)
+logger = logging.getLogger()
+
+
 
 batch_size = 128
 num_classes = 10
 dense_size = 16
-epochs = 4000
+epochs = 40
 
-max_epochs = 40000
+max_epochs = 400
 
 # input image dimensions
 img_rows, img_cols = 28, 28
@@ -41,8 +45,8 @@
 print(x_test.shape[0], "test samples")
 
 # convert class vectors to binary class matrices
-y_train = keras.utils.to_categorical(y_train, num_classes)
-y_test = keras.utils.to_categorical(y_test, num_classes)
+y_train = tf.keras.utils.to_categorical(y_train, num_classes)
+y_test = tf.keras.utils.to_categorical(y_test, num_classes)
 
 inputs = Input(shape=(28, 28, 1))
 flatten = Flatten()(inputs)
@@ -56,7 +60,7 @@
     myopt = evolutionary_keras.optimizers.CMA(population_size=5, sigma_init=15)
     epochs = 1
 else:
-    myopt = evolutionary_keras.optimizers.NGA(population_size=20, sigma_init=15)
+    myopt = evolutionary_keras.optimizers.NGA(population_size=2, sigma_init=15)
 
 print(" > Compiling the model")
 model.compile(optimizer=myopt, loss="categorical_crossentropy", metrics=["accuracy"])
@@ -71,6 +75,7 @@
     verbose=1,
     # validation_data=(x_test, y_test)
 )
-score = model.evaluate(x=x_test, y=y_test, verbose=0)
-print("Test loss:", score[0])
-print("Test accuracy:", score[1])
+score = model.evaluate(x=x_test, y=y_test, return_dict=True, verbose=0)
+
+print("Test loss:", score['loss'])
+print("Test accuracy:", score['accuracy'])
diff --git a/setup.py b/setup.py
index 46d125c..e860123 100644
--- a/setup.py
+++ b/setup.py
@@ -1,19 +1,44 @@
 from setuptools import setup, find_packages
-from os import path
+import os
+import re
 
-this_directory = path.abspath(path.dirname(__file__))
-with open(path.join(this_directory, "README.md"), encoding="utf-8") as f:
+
+requirements = ['numpy', 'cma']
+if os.environ.get('READTHEDOCS') != 'True':
+    requirements.append('tensorflow>2.1*')
+PACKAGE = 'evolutionary_keras'
+
+this_directory = os.path.abspath(os.path.dirname(__file__))
+with open(os.path.join(this_directory, "README.md"), encoding="utf-8") as f:
     long_description = f.read()
 
+def get_version():
+    """ Gets the version from the package's __init__ file
+    if there is some problem, let it happily fail """
+    VERSIONFILE = os.path.join('src', PACKAGE, '__init__.py')
+    initfile_lines = open(VERSIONFILE, 'rt').readlines()
+    VSRE = r"^__version__ = ['\"]([^'\"]*)['\"]"
+    for line in initfile_lines:
+        mo = re.search(VSRE, line, re.M)
+        if mo:
+            return mo.group(1)
+
 setup(
     name="evolutionary_keras",
-    version="1.0.1",
+    version=get_version(),
     author="S. Carrazza, J. Cruz-Martinez, Roy Stegeman",
     author_email="juan.cruz@mi.infn.it, roy.stegeman@mi.infn.it",
     url="https://github.com/N3PDF/evolutionary_keras",
     package_dir={"": "src"},
     packages=find_packages("src"),
-    install_requires=["numpy", "keras", "sphinx_rtd_theme", "recommonmark", "cma"],
+    zip_safe=False,
+    install_requires=requirements,
+    extras_require={
+        'docs' : [
+        'sphinx_rtd_theme',
+        'recommonmark',
+        ],
+    },
     python_requires=">=3.6",
     descriptions="An evolutionary algorithm implementation for Keras",
     long_description=long_description,
diff --git a/src/evolutionary_keras/__init__.py b/src/evolutionary_keras/__init__.py
index 5becc17..44759f5 100644
--- a/src/evolutionary_keras/__init__.py
+++ b/src/evolutionary_keras/__init__.py
@@ -1 +1 @@
-__version__ = "1.0.0"
+__version__ = "1.9"
diff --git a/src/evolutionary_keras/models.py b/src/evolutionary_keras/models.py
index 8f804ef..09058c4 100644
--- a/src/evolutionary_keras/models.py
+++ b/src/evolutionary_keras/models.py
@@ -2,11 +2,10 @@
 
 import logging
 
-from keras.callbacks.callbacks import History
-from keras.models import Model
+from tensorflow.keras.callbacks import History
+from tensorflow.keras.models import Model
 
 import evolutionary_keras.optimizers as Evolutionary_Optimizers
-from evolutionary_keras.utilities import parse_eval
 
 log = logging.getLogger(__name__)
 
@@ -27,7 +26,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.is_genetic = False
         self.opt_instance = None
-        self.history = History()
+        self.history_info = History()
 
     def parse_optimizer(self, optimizer):
         """ Checks whether the optimizer is genetic and creates and optimizer instance in case a
@@ -35,7 +34,8 @@ def parse_optimizer(self, optimizer):
         """
         # Checks (if the optimizer input is a string) and whether it is in the 'optimizers'
         # dictionary
-        if isinstance(optimizer, str) and optimizer in optimizer_dict.keys():
+
+        if isinstance(optimizer, str) and optimizer.lower() in optimizer_dict.keys():
             opt = optimizer_dict.get(optimizer.lower())
             # And instanciate it with default values
             optimizer = opt()
@@ -50,6 +50,7 @@ def compile(self, optimizer="rmsprop", **kwargs):
         """ When the optimizer is genetic, compiles the model in keras setting an arbitrary
         keras supported optimizer """
         self.parse_optimizer(optimizer)
+        self.history_info.set_model(self)
         if self.is_genetic:
             super().compile(optimizer="rmsprop", **kwargs)
         else:
@@ -71,7 +72,7 @@ def perform_genetic_fit(
                 verbose, prints to log.info the loss per epoch
         """
         # Prepare the history for the initial epoch
-        self.history.on_train_begin()
+        self.history_info.on_train_begin()
         # Validation data is currently not being used!!
         if validation_data is not None:
             log.warning(
@@ -84,29 +85,24 @@ def perform_genetic_fit(
                 "The optimizer determines the number of generations, epochs will be ignored."
             )
 
-        metricas = self.metrics_names
         for epoch in range(epochs):
             # Generate the best mutant
             score, best_mutant = self.opt_instance.run_step(x=x, y=y)
 
+            training_metric = next(iter(score))
+
             # Ensure the best mutant is the current one
             self.set_weights(best_mutant)
             if verbose == 1:
-                loss = parse_eval(score)
+                loss = score[training_metric]
                 information = f" > epoch: {epoch+1}/{epochs}, {loss} "
                 log.info(information)
+
             # Fill keras history
-            try:
-                history_data = dict(zip(metricas, score))
-            except TypeError as e:
-                # Maybe score was just one number, evil Keras
-                if parse_eval(score) == score:
-                    score = [score, score]
-                    history_data = dict(zip(metricas, score))
-                else:
-                    raise e
-            self.history.on_epoch_end(epoch, history_data)
-        return self.history
+            history_data = score
+            self.history_info.on_epoch_end(epoch, history_data)
+
+        return self.history_info
 
     def fit(self, x=None, y=None, validation_data=None, epochs=1, verbose=0, **kwargs):
         """ If the optimizer is genetic, the fitting procedure consists on executing `run_step` for
diff --git a/src/evolutionary_keras/optimizers.py b/src/evolutionary_keras/optimizers.py
index a8a8eca..65d763b 100644
--- a/src/evolutionary_keras/optimizers.py
+++ b/src/evolutionary_keras/optimizers.py
@@ -7,23 +7,19 @@
 
 import cma
 import numpy as np
-from keras.optimizers import Optimizer
+from tensorflow.keras.optimizers import Optimizer
 
-from evolutionary_keras.utilities import (
-    compatibility_numpy,
-    get_number_nodes,
-    parse_eval,
-)
+from evolutionary_keras.utilities import compatibility_numpy, get_number_nodes
 
 
 class EvolutionaryStrategies(Optimizer):
     """ Parent class for all Evolutionary Strategies
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, name, **kwargs):
+        super().__init__(name, **kwargs)
         self.model = None
-        self.shape = None
+        self.shape = []
         self.non_training_weights = []
 
     @abstractmethod
@@ -51,7 +47,13 @@ def run_step(self, x, y):
 
     def get_updates(self, loss, params):
         """ Capture Keras get_updates method """
-        pass
+
+    def _resource_apply_dense(self):
+        """ Override """
+
+    def _resource_apply_sparse(self):
+        """ Override """
+
 
 
 class NGA(EvolutionaryStrategies):
@@ -74,7 +76,12 @@ class NGA(EvolutionaryStrategies):
     # In case the user wants to adjust sigma_init
     # population_size or mutation_rate parameters the NGA method has to be initiated
     def __init__(
-        self, sigma_init=15, population_size=80, mutation_rate=0.05, *args, **kwargs
+        self,
+        name="NGA",
+        sigma_init=15,
+        population_size=80,
+        mutation_rate=0.05,
+        **kwargs
     ):
         self.sigma_init = sigma_init
         self.population_size = population_size
@@ -83,7 +90,18 @@ def __init__(
         self.n_nodes = 0
         self.n_generations = 1
 
-        super(NGA, self).__init__(*args, **kwargs)
+        super(NGA, self).__init__(name, **kwargs)
+
+    def get_config(self):
+        config = super(NGA, self).get_config()
+        config.update(
+            {
+                "sigma_init": self.sigma_init,
+                "population_size": self.population_size,
+                "mutation_rate": self.mutation_rate,
+            }
+        )
+        return config
 
     # Only works if all non_trainable_weights come after all trainable_weights
     # perhaps part of the functionality (getting shape) can be moved to ES
@@ -179,8 +197,9 @@ def evaluate_mutants(self, mutants, x=None, y=None, verbose=0):
             `loss`: loss of the best performing mutant
             `best_mutant`: best performing mutant
         """
-        best_loss = self.model.evaluate(x=x, y=y, verbose=verbose)
-        best_loss_val = parse_eval(best_loss)
+        best_loss = self.model.evaluate(x=x, y=y, verbose=verbose, return_dict=True)
+        training_metric = next(iter(best_loss))
+        best_loss_val = best_loss[training_metric]
         best_mutant = mutants[0]
         new_mutant = False
         for mutant in mutants[1:]:
@@ -188,8 +207,8 @@ def evaluate_mutants(self, mutants, x=None, y=None, verbose=0):
             # TODO related to the other todos, eventually this will have to be done
             # in a per-layer basis
             self.model.set_weights(mutant)
-            loss = self.model.evaluate(x=x, y=y, verbose=False)
-            loss_val = parse_eval(loss)
+            loss = self.model.evaluate(x=x, y=y, verbose=False, return_dict=True)
+            loss_val = loss[training_metric]
 
             if loss_val < best_loss_val:
                 best_loss_val = loss_val
@@ -237,12 +256,12 @@ class CMA(EvolutionaryStrategies):
 
     def __init__(
         self,
+        name="CMA",
         sigma_init=0.3,
         target_value=None,
         population_size=None,
         max_evaluations=None,
         verbosity=1,
-        *args,
         **kwargs
     ):
         """
@@ -251,7 +270,6 @@ def __init__(
         The default `epochs` in EvolModel is 1, meaning `run step` called once during training.
         """
         self.sigma_init = sigma_init
-        self.shape = None
         self.length_flat_layer = None
         self.trainable_weights_names = None
         self.verbosity = verbosity
@@ -268,7 +286,19 @@ def __init__(
         if population_size:
             self.options["popsize"] = population_size
 
-        super(CMA, self).__init__(*args, **kwargs)
+        super(CMA, self).__init__(name, **kwargs)
+
+    def get_config(self):
+        config = super(CMA, self).get_config()
+        config.update(
+            {
+                "sigma_init": self.sigma_init,
+                "target_value": self.target_value,
+                "population_size": self.population_size,
+                "max_evaluations": self.max_evaluations,
+            }
+        )
+        return config
 
     def on_compile(self, model):
         """ Function to be called by the model during compile time. Register the model `model` with
@@ -355,13 +385,16 @@ def run_step(self, x, y):
         # If max_evaluations is not set manually, use the number advised in arXiv:1604.00772
         if self.max_evaluations is None:
             self.options["maxfevals"] = 1e3 * len(x0) ** 2
+        else:
+            self.options["maxfevals"] = self.max_evaluations
 
         # minimizethis is function that 'cma' aims to minimize
         def minimizethis(flattened_weights):
             weights = self.undo_flatten(flattened_weights)
             self.model.set_weights(weights)
-            loss = parse_eval(self.model.evaluate(x=x, y=y, verbose=0))
-            return loss
+            loss = self.model.evaluate(x=x, y=y, verbose=0, return_dict=True)
+            training_metric = next(iter(loss))
+            return loss[training_metric]
 
         # Run the minimization and return the ultimatly selected 1 dimensional layer of weights
         # 'xopt'.
@@ -376,5 +409,5 @@ def minimizethis(flattened_weights):
 
         # Determine the ultimatly selected mutants' performance on the training data.
         self.model.set_weights(selected_parent)
-        loss = self.model.evaluate(x=x, y=y, verbose=0)
+        loss = self.model.evaluate(x=x, y=y, verbose=0, return_dict=True)
         return loss, selected_parent
diff --git a/src/evolutionary_keras/tests/test_fit.py b/src/evolutionary_keras/tests/test_fit.py
index 54f9714..a44f5f7 100644
--- a/src/evolutionary_keras/tests/test_fit.py
+++ b/src/evolutionary_keras/tests/test_fit.py
@@ -1,11 +1,11 @@
-""" Tests for checking that we can compile and run with all the included optimziers """
+""" Tests for checking that we can compile and run with all the included optimizers """
 
 import os
+
 # ensure the tests run on CPU
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 import numpy as np
-import keras.backend as K
-from keras.layers import Input, Dense
+from tensorflow.keras.layers import Dense, Input
 
 from evolutionary_keras.models import EvolModel
 from evolutionary_keras.optimizers import NGA, CMA
@@ -22,29 +22,30 @@ def generate_model(ishape=2, hshape=4, oshape=1):
 
 
 def test_NGA():
-    """ Test the NGA is able to run fit """
-    optimizer = NGA(mutation_rate=1.0)
+    """ Tests whether the NGA is able to run fit """
+    optimizer = NGA(mutation_rate=1.0, population_size=10)
     modelito = generate_model(ishape=2, hshape=4, oshape=1)
     modelito.compile(optimizer=optimizer, loss="mse")
     # Since we just want to check the optimizer is able to run
-    # Give some arbitrary input and output for 50 epochs
+    # Give some arbitrary input and output for 10 epochs
     xin = np.ones((100, 2))
     yout = np.zeros((100, 1))
     start_loss = modelito.evaluate(x=xin, y=yout)
-    _ = modelito.fit(x=xin, y=yout, epochs=50)
+    _ = modelito.fit(x=xin, y=yout, epochs=10)
     final_loss = modelito.evaluate(x=xin, y=yout)
     assert final_loss < start_loss
 
+
 def test_CMA():
-    """ Test the CMA implementation """
-    optimizer = CMA()
+    """ Tests whether the CMA is able to run fit """
+    optimizer = CMA(max_evaluations=100)
     modelito = generate_model(ishape=2, hshape=4, oshape=1)
     modelito.compile(optimizer=optimizer, loss="mse")
     # Since we just want to check the optimizer is able to run
-    # Give some arbitrary input and output for 50 epochs
+    # Give some arbitrary input and output for 100 epochs
     xin = np.ones((100, 2))
     yout = np.zeros((100, 1))
     start_loss = modelito.evaluate(x=xin, y=yout)
-    _ = modelito.fit(x=xin, y=yout, epochs=50)
+    _ = modelito.fit(x=xin, y=yout)
     final_loss = modelito.evaluate(x=xin, y=yout)
     assert final_loss < start_loss
diff --git a/src/evolutionary_keras/tests/test_utilities.py b/src/evolutionary_keras/tests/test_utilities.py
index ee1f728..53c4e08 100644
--- a/src/evolutionary_keras/tests/test_utilities.py
+++ b/src/evolutionary_keras/tests/test_utilities.py
@@ -1,27 +1,20 @@
 """ Test the utilities of evolutionary_keras """
 
-import numpy as np
-from keras import backend as K
-from keras.layers import Dense
-from evolutionary_keras.utilities import parse_eval, get_number_nodes
+from tensorflow.keras.layers import Dense, Input
 
-
-def test_parse_eval():
-    """ Test the parse_eval function, which should output a float
-    when it gets a list or a float """
-    # There are two situations, get a list or get a float
-    float_mode = 3.0
-    list_mode = [float_mode]
-    assert float_mode == parse_eval(list_mode)
-    assert float_mode == parse_eval(float_mode)
+from evolutionary_keras.utilities import get_number_nodes
+from evolutionary_keras.models import EvolModel
 
 
 def test_get_number_nodes():
     """ Get the number of nodes of a layer """
-    ii = K.constant(np.random.rand(2, 1))
+
     nodes = 10
-    layer = Dense(units=nodes)
-    # Keras won't build the layer until it is called with some input
-    _ = layer(ii)
+
+    # Tensorflow won't build the layer until it is called in a model
+    input_layer = Input(shape=(1,))
+    output_layer = Dense(units=nodes, name="test_layer")
+    modelito = EvolModel(input_layer, output_layer(input_layer))
+
     # Check that indeed the number of nodes is parsed correctly
-    assert nodes == get_number_nodes(layer)
+    assert nodes == get_number_nodes(output_layer)
diff --git a/src/evolutionary_keras/utilities.py b/src/evolutionary_keras/utilities.py
index 8205b43..72024d3 100644
--- a/src/evolutionary_keras/utilities.py
+++ b/src/evolutionary_keras/utilities.py
@@ -1,7 +1,7 @@
 """
     Module including some useful functions
 """
-from keras import backend as K
+from tensorflow.keras import backend as K
 
 
 def get_number_nodes(layer):
@@ -17,23 +17,8 @@ def get_number_nodes(layer):
     return nodes
 
 
-def parse_eval(loss):
-    """ Parses the result from `model.evaluate`, which sometimes
-    comes as a list and sometimes comes as one single float
-    Returns
-    -------
-        `loss` : loss as a float
-    """
-    try:
-        loss = loss[0]
-    except TypeError as e:
-        # If the output was a number then it is ok
-        if not isinstance(loss, (float, int)):
-            raise e
-    return loss
-
-
 def compatibility_numpy(weight):
+    """ Wrapper in case the evaluated keras object doesn't have a numpy() method """
     try:
         result = weight.numpy()
     except NotImplementedError: