Repository: giotto-ai/giotto-deep (language: python)

Issue:
Improve docstrings
**Is your feature request related to a problem? Please describe.**
<!--
A clear and concise description of what the problem is. Eg. "I'm always frustrated when [...]"
-->

Need better documentation to allow users to understand the code
**Describe the solution you'd like**
<!--
A clear and concise description of what you want to happen. If you have ideas on how to implement
your solution with (pseudo-)code, include them here.
-->
We need to:
 - [ ] improve the documentation, in particular the descriptions of the methods and the parameters
 - [ ] fix typos
 - [ ] add copy/paste code, like [here](https://github.com/giotto-ai/giotto-deep/blob/f62b2efcf2b42a88a8b86517024cff5b32a4bbd1/gdeep/pipeline/pipeline.py#L75)

**Describe alternatives you've considered**
<!--
A clear and concise description of any alternative solutions or features you've considered.
-->

**Additional context**
<!--
Add any other context or screenshots about the feature request here.
-->

<!-- Thanks for contributing! -->



Discussion / hints:

Fix patch:
diff --git a/docs/source/modules/analysis.rst b/docs/source/modules/analysis.rst
index 22ca7087..b7078f61 100644
--- a/docs/source/modules/analysis.rst
+++ b/docs/source/modules/analysis.rst
@@ -1,8 +1,14 @@
 Analysis
 ========
 
+Interpretability
+----------------
+
 .. automodule:: gdeep.analysis.interpretability
    :members:
-   
+
+Decision Boundary
+-----------------
+
 .. automodule:: gdeep.analysis.decision_boundary
    :members:
diff --git a/docs/source/modules/utility.rst b/docs/source/modules/utility.rst
index dd4a6f4f..da05cd0d 100644
--- a/docs/source/modules/utility.rst
+++ b/docs/source/modules/utility.rst
@@ -4,3 +4,8 @@ Utility
 .. automodule:: gdeep.utility
    :members:
    
+Persistence Gradient
+--------------------
+
+.. automodule:: gdeep.utility.optimisation
+   :members:
\ No newline at end of file
diff --git a/examples/Ensemble_learning.ipynb b/examples/basic_tutorial_ensemble_learning.ipynb
similarity index 74%
rename from examples/Ensemble_learning.ipynb
rename to examples/basic_tutorial_ensemble_learning.ipynb
index 4fff4ad8..25bec411 100644
--- a/examples/Ensemble_learning.ipynb
+++ b/examples/basic_tutorial_ensemble_learning.ipynb
@@ -8,9 +8,18 @@
     "# Ensemble learning\n",
     "#### Author: Matteo Caorsi\n",
     "\n",
-    "It is possible, in giotto-deep, to use the ensemble models of the library `ensemble-pxytorch` together with all the functionalities of `gdeep`.\n",
+    "It is possible, in giotto-deep, to use the ensemble models of the library `ensemble-pytorch` together with all the functionalities of `gdeep`.\n",
     "\n",
-    "In this short tutorial we will explain this concept with a simple example."
+    "## Scope\n",
+    "\n",
+    "Ensamble technique put together the predictions of different models and decide which is the best answer. It is a bit like having and ensemble of experts giving opinions and then the person in charge takes the final decision. In this example, we will try out the `VotingClassifer`, i.e. an ensemble method that decides on the best preediciton based on the majority of experts votes.\n",
+    "\n",
+    "## Content\n",
+    "These aree the main steps we wll follow:\n",
+    " 1. Load your data\n",
+    " 2. Defne a single expert\n",
+    " 3. wrap the ensemble model\n",
+    " 4. train the ensemble"
    ]
   },
   {
@@ -23,14 +32,16 @@
     "# imports\n",
     "from torch import nn\n",
     "from torch.utils.tensorboard import SummaryWriter\n",
-    "from gdeep.models import FFNet\n",
+    "from torch.optim.lr_scheduler import ExponentialLR\n",
+    "from torch.optim import SGD, Adam\n",
     "from torchensemble import VotingClassifier\n",
+    "\n",
     "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
     "from gdeep.trainer import Trainer\n",
     "from gdeep.utility.optimisation import SAMOptimizer\n",
-    "from torch.optim.lr_scheduler import ExponentialLR\n",
-    "from torch.optim import SGD, Adam\n",
-    "from gdeep.utility import ensemble_wrapper"
+    "from gdeep.models import FFNet\n",
+    "from gdeep.utility import ensemble_wrapper\n",
+    "from gdeep.visualisation import Visualiser"
    ]
   },
   {
@@ -40,8 +51,8 @@
    "source": [
     "# Initialize the tensorboard writer\n",
     "\n",
-    "In order to analyse the reuslts of your models, you need to start tensorboard.\n",
-    "On the terminal, move inside the `/example` folder. There run the following command:\n",
+    "In order to analyse the results of your models, you need to start tensorboard.\n",
+    "On the terminal, move inside the `/examples` folder. There run the following command:\n",
     "\n",
     "```\n",
     "tensorboard --logdir=runs\n",
@@ -65,7 +76,9 @@
    "id": "a956ff40-cc51-49bb-bd72-34a05f06ea99",
    "metadata": {},
    "source": [
-    "# Load your data"
+    "# Load your data\n",
+    "\n",
+    "In this example we use a tabular dataset and the task is a classification task. The dataset is a point cloud representing two entangled tori and the model needs to classify each point as belonging to one or the other torus."
    ]
   },
   {
@@ -89,7 +102,7 @@
    "source": [
     "# Define a single estimator of the ensemble\n",
     "\n",
-    "You can define it as you would normally do with any neural network"
+    "You can define a single estimator of an ensemble as you would normally do with any other neural network"
    ]
   },
   {
@@ -139,7 +152,7 @@
    "id": "ecd8167e-9b7b-43a4-9b2b-9b28b303def4",
    "metadata": {},
    "source": [
-    "What you would have done instead:\n",
+    "What you would have done instead, renouncing to many giotto-deep capablities:\n",
     "\n",
     "```\n",
     "model = VotingClassifier(\n",
@@ -148,7 +161,7 @@
     "    cuda=False\n",
     ")\n",
     "model.set_optimizer(\"Adam\")\n",
-    "model.fit(train_loader=dl_tr,epochs=1`\n",
+    "model.fit(train_loader=dl_tr,epochs=1)\n",
     "```\n"
    ]
   },
@@ -157,8 +170,8 @@
    "id": "8f6d96a5-8de4-4d2d-8048-dcc0e8a2bed2",
    "metadata": {},
    "source": [
-    "# Train your model\n",
-    "You can easily train your esnsemble model as you would train any other model in giotto-deep."
+    "# Train your ensemble of models\n",
+    "You can easily train your ensemble model as you would train any other model in giotto-deep: initialise the `Trainer` class and run it."
    ]
   },
   {
@@ -188,7 +201,7 @@
    "source": [
     "# Visualise the model graph\n",
     "\n",
-    "You can integractively visualise your model by checking it on tensorboard after these few lines:"
+    "You can integractively visualise your ensemble of models by checking it on tensorboard after these few lines are executed:"
    ]
   },
   {
@@ -198,11 +211,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.visualisation import Visualiser\n",
-    "\n",
+    "# initialise he visualiser\n",
     "vs = Visualiser(pipe)\n",
     "\n",
-    "vs.plot_data_model()  # send the graph of the model to tb"
+    "# send the graph of the model to tensorboard\n",
+    "vs.plot_data_model()  "
    ]
   },
   {
diff --git a/examples/basic_tutorial_gridsearch_benchmarking.ipynb b/examples/basic_tutorial_hpo_and_benchmarking.ipynb
similarity index 77%
rename from examples/basic_tutorial_gridsearch_benchmarking.ipynb
rename to examples/basic_tutorial_hpo_and_benchmarking.ipynb
index 02bba6af..47d32566 100644
--- a/examples/basic_tutorial_gridsearch_benchmarking.ipynb
+++ b/examples/basic_tutorial_hpo_and_benchmarking.ipynb
@@ -4,17 +4,23 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Basic tutorial: gridsearch and benchmarking\n",
+    "# Basic tutorial: HPO and benchmarking\n",
     "#### Author: Matteo Caorsi\n",
     "\n",
-    "This short tutorial provides you with the basic functioning of *giotto-deep* API.\n",
+    "This tutorial is focussed around **hyperparameter optimisation** (HPO), a rather unique feature of giotto-deep compared to other deep-learning frameworks.\n",
+    "\n",
+    "## Scope\n",
+    "\n",
+    "Neural network are very complex beasts and it is not at all intuitive to understand whether the change in some of the structural parameters would lead to an improved performance and higher stability. For example, is it always better to increase the depth of a feed-forward network? The answer is, in general, \"no\": it really depends on the complexity of your problem (*VC dimension and friends*). However, given that it is many time impossible to compute a priori quantities like the [VC dimension](https://en.wikipedia.org/wiki/Vapnik–Chervonenkis_dimension), it is more effective to empirically try out different structural parameters and benchmark the results. This is what HPO is all about.\n",
+    "\n",
+    "## Plan for the tutorial\n",
     "\n",
     "The main steps of the tutorial are the following:\n",
     " 1. creation of a dataset\n",
     " 2. creation of a model\n",
-    " 3. gridsearch on the initial model and dataset\n",
+    " 3. HPO on the initial model and dataset\n",
     " 4. benchmarking\n",
-    " 5. gridsearching in each benchmark"
+    " 5. HPO-ing in each benchmark"
    ]
   },
   {
@@ -29,33 +35,26 @@
     "%matplotlib inline\n",
     "\n",
     "import numpy as np\n",
-    "\n",
     "import torch\n",
     "from torch import nn\n",
-    "\n",
-    "from gdeep.models import FFNet\n",
-    "\n",
-    "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
-    "\n",
-    "from torch.utils.tensorboard import SummaryWriter\n",
+    "from torch.utils.tensorboard.writer import SummaryWriter\n",
     "from torch.utils.data.sampler import SubsetRandomSampler\n",
     "import torchvision.models as models\n",
-    "\n",
+    "from gtda.diagrams import BettiCurve\n",
+    "from gtda.plotting import plot_betti_surfaces\n",
+    "from sklearn.model_selection import StratifiedKFold, KFold\n",
+    "import optuna\n",
     "from torch.optim import SGD, Adam, RMSprop\n",
     "from torch.optim.lr_scheduler import ExponentialLR\n",
+    "\n",
+    "from gdeep.data.preprocessors import ToTensorImage\n",
     "from gdeep.trainer import Trainer\n",
     "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
     "# today's protagonists\n",
     "from gdeep.search import Benchmark\n",
     "from gdeep.search import HyperParameterOptimization, GiottoSummaryWriter\n",
-    "\n",
-    "from gtda.diagrams import BettiCurve\n",
-    "\n",
-    "from gtda.plotting import plot_betti_surfaces\n",
-    "\n",
-    "from sklearn.model_selection import StratifiedKFold, KFold\n",
-    "\n",
-    "import optuna"
+    "from gdeep.models import FFNet\n",
+    "from gdeep.visualisation import  persistence_diagrams_of_activations\n"
    ]
   },
   {
@@ -91,7 +90,7 @@
    "source": [
     "# Create your dataset\n",
     "\n",
-    "In the next cell we subsample the `CIFAR10` dataset and prepare the data loaders"
+    "In the next cell we subsample the [CIFAR10](https://www.cs.toronto.edu/~kriz/cifar.html) dataset and prepare the data loaders. Note that a preprocessing step is required to transformed the images into tensors."
    ]
   },
   {
@@ -100,15 +99,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from torchvision.transforms import ToTensor\n",
     "\n",
+    "# download the dataset\n",
     "bd = DatasetBuilder(name=\"CIFAR10\")\n",
-    "train_indices = list(range(32*10))\n",
-    "\n",
-    "print(len(train_indices))\n",
-    "\n",
-    "ds_tr, _, _ = bd.build( download=True)\n",
-    "\n"
+    "ds_tr, _, _ = bd.build( download=True)\n"
    ]
   },
   {
@@ -117,9 +111,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Preprocessing\n",
-    "\n",
-    "from gdeep.data.preprocessors import ToTensorImage\n",
+    "# Preprocessing steps\n",
     "\n",
     "transformation = ToTensorImage((32,32))\n",
     "transformation.fit_to_dataset(ds_tr)  # this is useless for this transformation\n",
@@ -138,7 +130,7 @@
    "source": [
     "## Define your model\n",
     "\n",
-    "In the next section we build a torchh model with a `str` parameter. The type of parameter can of course be changed to `int`: the example is to show the potential of the code."
+    "In the next section we build a torch model with a `str` parameter. The type of parameter can of course be changed to `int`: the example is to show the potential of the code."
    ]
   },
   {
@@ -152,7 +144,7 @@
     "    def __init__(self, n_nodes = \"100\"):\n",
     "        super(model2, self).__init__()\n",
     "        self.md = nn.Sequential(nn.Sequential(models.resnet18(pretrained=True), \n",
-    "                                                  nn.Linear(1000,eval(n_nodes))), \n",
+    "                                                  nn.Linear(1000, eval(n_nodes))), \n",
     "                                    nn.Linear(eval(n_nodes),10))\n",
     "    \n",
     "    def forward(self, x):\n",
@@ -162,6 +154,15 @@
     "model = model2()\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Training without HPO\n",
+    "\n",
+    "This step is the normal, non HPO, step that you would do to train your model. Starting from the next section on, we will dive into the HPO framework."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -172,7 +173,7 @@
     "loss_fn = nn.CrossEntropyLoss()\n",
     "\n",
     "# initialise pipeline class\n",
-    "pipe = Trainer(model, [dl_tr, None], loss_fn, writer, StratifiedKFold(2, shuffle=True))"
+    "pipe = Trainer(model, [dl_tr, None], loss_fn, writer, k_fold_class=StratifiedKFold(2, shuffle=True))"
    ]
   },
   {
@@ -181,8 +182,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# the following is a simple cross-validated training (no Gridsearch)\n",
-    "\n",
+    "# the following is a simple cross-validated training (no HPO)\n",
     "# we also add the n_accumulated_grads=5, which is useful to avoid OOM results when training on the GPU\n",
     "pipe.train(SGD, 2, True, {\"lr\": 0.001}, n_accumulated_grads=5)"
    ]
@@ -193,12 +193,13 @@
    "source": [
     "# HyperParameter Optimization\n",
     "\n",
+    "One of the unique features of gotto-deep is the possibility to run advanced hyperparameters searches in a few lines of code: it is enough to define the dictionaries of hyperparameters, initialise the class `HyperParameterOptimization` and run it with the method `start`.\n",
     "We run a search over different hyperparametrs: \n",
     " - the learning rate `lr`\n",
     " - the batch size `batch_size`\n",
     " - the network parameter `arch`\n",
     " \n",
-    "The scope of the search is to find the optimum set of hyperparameters."
+    "The scope of the search is to find the optimum set of hyperparameters. The \"optimum\" depends on either the accuracy (or the user-defined metric) or the loss."
    ]
   },
   {
@@ -219,7 +220,7 @@
     "dataloaders_params = {\"batch_size\": [32, 64, 16]}\n",
     "models_hyperparams = {\"n_nodes\": [\"200\"]}\n",
     "\n",
-    "# starting the gridsearch\n",
+    "# starting the HPO\n",
     "search.start((SGD, Adam), 3, False, optimizers_params, dataloaders_params, models_hyperparams, n_accumulated_grads=2)"
    ]
   },
@@ -229,7 +230,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(search.best_val_acc_gs, search.best_val_loss_gs)"
+    "print(\"These are the best results we have found so far: \", search.best_val_acc_gs, search.best_val_loss_gs)"
    ]
   },
   {
@@ -273,7 +274,7 @@
    "source": [
     "# Benchmarking\n",
     "\n",
-    "Benchmarking means fixing a set of models and a set of datasets and trying all possible pairs of *(model, dataset)*. \n",
+    "Benchmarking means fixing a set of models and a set of datasets and trying all possible pairs of *(model, dataset)*. The most common usecase is actually to also fix the model and to run it over many datasets.\n",
     "\n",
     "Of course, only compatible models with compatiible datasets will be benchmarked.\n",
     "\n",
@@ -376,7 +377,7 @@
    "source": [
     "## Start the benchmarking!\n",
     "\n",
-    "after initialising th class with the dictionaries of models and dataloaders, we can run the actual benchmark."
+    "After initialising the class with the dictionaries of models and dataloaders, we can run the actual benchmark."
    ]
   },
   {
@@ -398,7 +399,7 @@
    "source": [
     "# Benchmarking + HyperParameter Optimization + CV\n",
     "\n",
-    "In this last section we consider the possibility of running a gridsearch within each pair *(model, dataset)*.\n",
+    "In this last section we consider the possibility of running an HPO within each pair *(model, dataset)*.\n",
     "\n",
     "This can be achieved by initialising a benchmark class and use the benchmark as input for the gridsearch class.\n",
     "\n",
@@ -415,7 +416,7 @@
     "loss_fn = nn.CrossEntropyLoss()\n",
     "\n",
     "# initialise benchmark\n",
-    "bench = Benchmark(models_dicts, dataloaders_dicts, loss_fn, writer, KFold(3))\n",
+    "bench = Benchmark(models_dicts, dataloaders_dicts, loss_fn, writer, k_fold_class=KFold(3))\n",
     "\n",
     "# initialise gridsearch with benchmark instance\n",
     "search2 = HyperParameterOptimization(bench, \"loss\", 2)\n",
@@ -426,7 +427,7 @@
     "models_hyperparams = {\"n_nodes\": [\"500\", \"200\"]}\n",
     "search2.start((SGD, Adam), 2, True, optimizers_params, dataloaders_params, models_hyperparams)\n",
     "\n",
-    "writer.close()  # let's recall to close the tensorboard writer once all is done"
+    "writer.close()  # let's not forget to close the tensorboard writer once all is done"
    ]
   },
   {
@@ -447,7 +448,9 @@
    "source": [
     "## Custom pruner and sampler\n",
     "\n",
-    "It is possible to pass to the HyperParameterOptimization class a customer `optuna.Pruners` and `optuna.Samplers`"
+    "It is possible to pass to the HyperParameterOptimization class a customer `optuna.Pruners` and `optuna.Samplers`.\n",
+    "\n",
+    "The pruner is used to stop a trial when it is clearly not rach acceptable performances (hence spearing a bit of computational costs), while a sampler is used to sample the hyperparameter space using different techniques."
    ]
   },
   {
@@ -459,18 +462,20 @@
     "from optuna.pruners import MedianPruner\n",
     "from optuna.samplers import TPESampler\n",
     "\n",
-    "# initialise gridsearch\n",
-    "gs = HyperParameterOptimization(pipe, \"accuracy\", 8, best_not_last=False, pruner=MedianPruner(n_startup_trials=2,\n",
-    "                                                                              n_warmup_steps=0,\n",
-    "                                                                              interval_steps=1,\n",
-    "                                                                              n_min_trials=1), sampler=TPESampler())\n",
+    "# initialise te HPO\n",
+    "gs = HyperParameterOptimization(pipe, \"accuracy\", 5, best_not_last=False, \n",
+    "                                pruner=MedianPruner(n_startup_trials=2,\n",
+    "                                                    n_warmup_steps=0,\n",
+    "                                                    interval_steps=1,\n",
+    "                                                    n_min_trials=1), \n",
+    "                                sampler=TPESampler())\n",
     "\n",
     "# dictionaries of hyperparameters\n",
     "optimizers_params = {\"lr\": [0.001, 0.01]}\n",
     "dataloaders_params = {\"batch_size\": [32, 64, 16]}\n",
     "models_hyperparams = {\"n_nodes\": [\"500\", \"200\"]}\n",
     "\n",
-    "# starting the gridsearch\n",
+    "# starting the HPO\n",
     "gs.start((SGD, Adam), 3, False, optimizers_params, dataloaders_params, models_hyperparams, n_accumulated_grads=2)"
    ]
   },
diff --git a/examples/basic_tutorial_image.ipynb b/examples/basic_tutorial_image.ipynb
index 2b348c94..cb94cd38 100644
--- a/examples/basic_tutorial_image.ipynb
+++ b/examples/basic_tutorial_image.ipynb
@@ -7,7 +7,13 @@
     "# Basic tutorial: image data\n",
     "#### Author: Matteo Caorsi\n",
     "\n",
-    "This short tutorial provides you with the basic functioning of *giotto-deep* API.\n",
+    "This short tutorial shows you how to use the basic functioning of *giotto-deep* API for an imagee classfication task.\n",
+    "\n",
+    "Image classifcation is about associating a picture with a label: for example, you would like your model to be able to associate a label `cat` to the image of a cat and a label `dog`to an image of a dog:\n",
+    "\n",
+    "![img](./images/image_class.png)\n",
+    "\n",
+    "It is important to clarify that the set of labels has to be pre-defined and the modls have to be trained with a sufficient amount of data per label.\n",
     "\n",
     "The main steps of the tutorial are the following:\n",
     " 1. creation of a dataset\n",
@@ -28,21 +34,19 @@
     "%matplotlib inline\n",
     "\n",
     "import numpy as np\n",
-    "\n",
     "import torch\n",
     "from torch import nn\n",
+    "from torch.utils.data.sampler import SubsetRandomSampler\n",
+    "from torch.utils.tensorboard.writer import SummaryWriter\n",
+    "import torchvision.models as models\n",
+    "from gtda.diagrams import BettiCurve\n",
+    "from gtda.plotting import plot_betti_surfaces\n",
     "\n",
+    "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
     "from gdeep.models import FFNet\n",
-    "\n",
     "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
-    "\n",
-    "from torch.utils.tensorboard import SummaryWriter\n",
-    "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
-    "\n",
-    "\n",
-    "from gtda.diagrams import BettiCurve\n",
-    "\n",
-    "from gtda.plotting import plot_betti_surfaces"
+    "from gdeep.data.preprocessors import ToTensorImage\n",
+    "from gdeep.trainer import Trainer"
    ]
   },
   {
@@ -51,14 +55,14 @@
    "source": [
     "# Initialize the tensorboard writer\n",
     "\n",
-    "In order to analyse the reuslts of your models, you need to start tensorboard.\n",
-    "On the terminal, move inside the `/example` folder. There run the following command:\n",
+    "In order to analyse the results of your models, you need to start tensorboard.\n",
+    "On the terminal, move inside the `/examples` folder. There run the following command:\n",
     "\n",
     "```\n",
     "tensorboard --logdir=runs\n",
     "```\n",
     "\n",
-    "Then go [here](http://localhost:6006/) after the training to see all the visualisation results."
+    "Then go [here](http://localhost:6006/) after the training of your model to see all the visualisation results."
    ]
   },
   {
@@ -74,7 +78,11 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Create your dataset"
+    "# Create your dataset\n",
+    "\n",
+    "In this example we will be using the `CIFAR10` dataset, which is a dataset of many thousands of images with ten classes.\n",
+    "\n",
+    "In giotto-deep it is enough to call `DatasetBuilder(name=\"CIFAR10\")` to get the dataset!"
    ]
   },
   {
@@ -83,11 +91,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from torch.utils.data.sampler import SubsetRandomSampler\n",
-    "from torchvision.transforms import ToTensor  # for preprocessing\n",
     "\n",
     "db = DatasetBuilder(name=\"CIFAR10\")\n",
-    "ds_tr, ds_val, ds_ts = db.build()\n"
+    "ds_tr, ds_val, ds_ts = db.build()\n",
+    "NUMBER_OF_CLASSES = 10"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Preprocessing\n",
+    "\n",
+    "The preprocessing steps consist in transformin the imagees to tensors. We have already provided a simple class to do such preprocessing: `ToTensorImage`.\n",
+    "\n",
+    "The transormation is computed on the fly, hence bring very light on the memory."
    ]
   },
   {
@@ -98,16 +116,16 @@
    "source": [
     "# Preprocessing\n",
     "\n",
-    "from gdeep.data.preprocessors import ToTensorImage\n",
+    "\n",
     "\n",
     "transformation = ToTensorImage((32,32))\n",
-    "transformation.fit_to_dataset(ds_tr)  # this is useless for this transformation\n",
+    "transformation.fit_to_dataset(ds_tr)  # this is useless for this transformation, but in general this is the API\n",
     "\n",
     "transformed_ds_tr = transformation.attach_transform_to_dataset(ds_tr)\n",
     "transformed_ds_val = transformation.attach_transform_to_dataset(ds_val)\n",
     "transformed_ds_ts = transformation.attach_transform_to_dataset(ds_ts)\n",
     "\n",
-    "# use only 320 images from cifar10\n",
+    "# use only 320 images from cifar10 for training\n",
     "train_indices = list(range(32*10))\n",
     "val_indices = list(range(32*5))\n",
     "test_indices = list(range(32*5))\n",
@@ -122,7 +140,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Define and train your model"
+    "## Define and train your model\n",
+    "\n",
+    "In the next cell we show how simple it is to take a model from `torch`, like [ResNet18](https://pytorch.org/vision/main/models/generated/torchvision.models.resnet18.html) and use it for training in Giotto-deep."
    ]
   },
   {
@@ -131,20 +151,31 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import torchvision.models as models\n",
-    "from gdeep.trainer import Trainer\n",
     "\n",
     "# wrap a sequential model in a torch nn.Module\n",
     "class model3(nn.Module):\n",
     "    def __init__(self):\n",
     "        super(model3, self).__init__()\n",
-    "        self.seqmodel = nn.Sequential(models.resnet18(pretrained=True), nn.Linear(1000,10))\n",
+    "        self.seqmodel = nn.Sequential(models.resnet18(pretrained=True), \n",
+    "                                      nn.Linear(1000, NUMBER_OF_CLASSES))\n",
+    "\n",
     "    def forward(self, X):\n",
     "        return self.seqmodel(X)\n",
     "\n",
     "model = model3()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Training the model\n",
+    "\n",
+    "We have prepard te dataset, we have prepared the model... it is now tim to train it!\n",
+    "\n",
+    "In giotto-deep it is just a matter of initialising the `Trainer` and run it with `train`. All the outputs, the exceptions, cross validation... everything you is just a parameter for the trainer. Have a look at the next cell!"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -153,7 +184,6 @@
    "source": [
     "from torch.optim import SGD\n",
     "\n",
-    "# print(model)\n",
     "loss_fn = nn.CrossEntropyLoss()\n",
     "\n",
     "pipe = Trainer(model, (dl_tr, dl_ts), loss_fn, writer)\n",
@@ -168,7 +198,33 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Simply use interpretability tools"
+    "# Evaluate model\n",
+    "\n",
+    "In the next section we compute the confusion matrix on the entire training dataloader."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pipe.evaluate_classification(NUMBER_OF_CLASSES)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Simply use interpretability tools\n",
+    "\n",
+    "After th training we would like to inspect what the model has learned. This is possible thanks to interpretability tools.\n",
+    "\n",
+    "Here below we show you how simple it is, in giotto-deep, to use many interpretabnility tools in one single line of code.\n",
+    "\n",
+    "The example below uses the `GuidedGradCAM` algorithm: given an input `datum`, the modl computes the gradients via backpropagation till the specified layder `conv1` and then averages over teh chanels of the convolutional layer.\n",
+    "\n",
+    "The result is then displayed as superimposition of the initial image, with a red/green hue to identify which pixes are pushing for and against the classification of teh given image into the selected `class`."
    ]
   },
   {
@@ -181,8 +237,10 @@
     "from gdeep.visualisation import Visualiser\n",
     "\n",
     "inter = Interpreter(pipe.model, method=\"GuidedGradCam\")\n",
-    "output = inter.interpret_image(next(iter(dl_tr))[0][0].reshape(1,3,32,32), \n",
-    "                      1, pipe.model.seqmodel[0].layer2[0].conv1);\n",
+    "datum = next(iter(dl_tr))[0][0].reshape(1,3,32,32)\n",
+    "layer = pipe.model.seqmodel[0].layer2[0].conv1\n",
+    "class_ = 1\n",
+    "output = inter.interpret_image(datum, class_, layer);\n",
     "\n",
     "# visualise the interpreter\n",
     "vs = Visualiser(pipe)\n",
@@ -196,7 +254,18 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Extract inner data from your models"
+    "# Appendix\n",
+    "\n",
+    "More advanced stuff ahead!\n",
+    "\n",
+    "## Extract inner data from your models\n",
+    "\n",
+    "For deeper analysis of the working of deep learning models, it is possible, in giotto-deep, to easily extract tons of parameters from your models:\n",
+    "Here below we show you how to extract:\n",
+    " - each layer parameter,\n",
+    " - the decision boundary,\n",
+    " - the activation functions given an input datum,\n",
+    " - the gradients at each node, given an input and a target."
    ]
   },
   {
@@ -209,10 +278,10 @@
     "\n",
     "me = ModelExtractor(pipe.model, loss_fn)\n",
     "\n",
-    "lista = me.get_layers_param()\n",
+    "list_of_layers = me.get_layers_param()\n",
     "\n",
-    "for k, item in lista.items():\n",
-    "    print(k,item.shape)\n"
+    "for k, item in list_of_layers.items():\n",
+    "    print(k, item.shape)\n"
    ]
   },
   {
@@ -221,8 +290,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# the decision boundary will be available on tensorboard, in the projectors section.\n",
     "x = next(iter(dl_tr))[0][0]\n",
-    "if x.dtype is not torch.int64:\n",
+    "if x.dtype is not torch.int64:  # cannot backpropagate on integers!\n",
     "    res = me.get_decision_boundary(x, n_epochs=1)\n",
     "    res.shape"
    ]
@@ -244,8 +314,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "x, target = next(iter(dl_tr))\n",
-    "if x.dtype is torch.float:\n",
+    "x, target = next(iter(dl_tr))  # a whole batch!\n",
+    "if x.dtype is torch.float:  # cannot backpropagate on integers!\n",
     "    for gradient in me.get_gradients(x, target=target)[1]:\n",
     "        print(gradient.shape)"
    ]
@@ -254,7 +324,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Visualise activations and other topological aspects of your model"
+    "## Visualise activations and other topological aspects of your model\n",
+    "\n",
+    "Finally, we show you how it is possible to extract, layer per layer, given a certain input batch, the persistence diagrams computed on the activation space. This means that, for a fixed layer, the input batch is representated as a point cloud in the space of activations for that selected layer. Out of this point cloud we compute the persistence diagrams. In giotto-deep this is literally one line of code."
    ]
   },
   {
@@ -264,29 +336,10 @@
    "outputs": [],
    "source": [
     "\n",
-    "# vs.plot_data_model()\n",
-    "# vs.plot_activations(x)\n",
+    "# the diagrams can be seen on tensorboard!\n",
     "vs.plot_persistence_diagrams(x)\n"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Evaluate model\n",
-    "\n",
-    "In the next section we compute the confusion matrix on the entire training dataloader."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "pipe.evaluate_classification(10)"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
diff --git a/examples/basic_tutorial_QandA.ipynb b/examples/basic_tutorial_question_answering.ipynb
similarity index 100%
rename from examples/basic_tutorial_QandA.ipynb
rename to examples/basic_tutorial_question_answering.ipynb
diff --git a/examples/basic_tutorial_regression.ipynb b/examples/basic_tutorial_regression.ipynb
index c3ea4be5..455267bf 100644
--- a/examples/basic_tutorial_regression.ipynb
+++ b/examples/basic_tutorial_regression.ipynb
@@ -7,7 +7,18 @@
     "# Basic tutorial: regression task\n",
     "#### Author: Matteo Caorsi\n",
     "\n",
-    "This short tutorial provides you with the basic functioning of *giotto-deep* API.\n",
+    "This short tutorial provides you with the basic example of regression.\n",
+    "Regression tasks can be performed in *giotto-deep* as easily as classification ones.\n",
+    "\n",
+    "## Scope\n",
+    "\n",
+    "The scope of a regression task in similar to a classificaiton one, with the difference that the space of classes is infinite! Hence, we cannot really ho to extend the losses and metrics of classificaiton tasks directly to regressions ones: however, the models, the training, th optimisation,... remain basically the same!\n",
+    "\n",
+    "![img](./images/regression.jpeg)\n",
+    "\n",
+    "In this tutorial we will try to fit a line in a 3D space: not super exciting indeed, but then you can have fun and use these tools to fit the curves in the *forex market* for example!\n",
+    "\n",
+    "## Content\n",
     "\n",
     "The main steps of the tutorial are the following:\n",
     " 1. creation of a dataset\n",
@@ -32,24 +43,17 @@
     "import torch\n",
     "from torch import nn\n",
     "import pandas as pd\n",
-    "\n",
     "from sklearn import datasets\n",
-    "\n",
-    "from gdeep.models import FFNet\n",
-    "from gdeep.models import ModelExtractor\n",
-    "from gdeep.analysis.interpretability import Interpreter\n",
+    "from gtda.diagrams import BettiCurve\n",
+    "from gtda.plotting import plot_betti_surfaces\n",
     "from torch.optim.lr_scheduler import ExponentialLR\n",
     "from torch.optim import SGD, Adam\n",
     "\n",
+    "from gdeep.models import FFNet\n",
     "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
     "from gdeep.trainer import Trainer\n",
-    "\n",
-    "from torch.utils.tensorboard import SummaryWriter\n",
     "from gdeep.data.datasets import DatasetBuilder, FromArray, DataLoaderBuilder\n",
-    "\n",
-    "from gtda.diagrams import BettiCurve\n",
-    "\n",
-    "from gtda.plotting import plot_betti_surfaces"
+    "from torch.utils.tensorboard import SummaryWriter\n"
    ]
   },
   {
@@ -58,14 +62,14 @@
    "source": [
     "# Initialize the tensorboard writer\n",
     "\n",
-    "In order to analyse the reuslts of your models, you need to start tensorboard.\n",
-    "On the terminal, move inside the `/example` folder. There run the following command:\n",
+    "In order to analyse the results of your models, you need to start tensorboard.\n",
+    "On the terminal, move inside the `/example` folder. There you can run the following command:\n",
     "\n",
     "```\n",
     "tensorboard --logdir=runs\n",
     "```\n",
     "\n",
-    "Then go [here](http://localhost:6006/) after the training to see all the visualisation results."
+    "Then go [here](http://localhost:6006/), after the training phase, to see all the visualisation results."
    ]
   },
   {
@@ -81,7 +85,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Create your dataset"
+    "# Create your dataset\n",
+    "\n",
+    "The dataset that we creeate is a 3D dataset representing a noisy hyperplane. The `X` is generated at random while the `y` has a linear relation eith `X`. The goal is to see if we can get `0.3` as linear coefficient once the modl has been trained."
    ]
   },
   {
@@ -105,7 +111,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Define and train your model"
+    "## Define and train your model\n",
+    "\n",
+    "The model is a simple feed-forward network: simple task, simple mdoel."
    ]
   },
   {
@@ -127,19 +135,12 @@
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [],
    "source": [
+    "## Define your metric\n",
     "\n",
-    "loss_fn = nn.MSELoss()\n",
-    "\n",
-    "pipe = Trainer(model, (dl_tr, dl_val), loss_fn, writer)\n",
-    "\n",
-    "# train the model with learning rate scheduler\n",
-    "pipe.train(Adam, 3, False, lr_scheduler=ExponentialLR, scheduler_params={\"gamma\": 0.9}, \n",
-    "           profiling=False, store_grad_layer_hist=True, writer_tag=\"line\")\n"
+    "In case of regression tasks, accuracy is not really a good metric. We propose here to compuite the $L_1$ norm between the prediction and the input"
    ]
   },
   {
@@ -148,16 +149,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "pipe.model(torch.tensor([[1,1,1]]).float())"
+    "def l1_norm(prediction, y):\n",
+    "    return torch.norm(prediction -y, p=1)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Let's train the model with cross validation: we just have to set the parameter `cross_validation = True`.\n",
+    "# Train the model\n",
     "\n",
-    "The `keep_training = True` flag allow us to restart from the same scheduler, optimiser and trained model obtained at thhe end of the last training in the instance of the class `pipe`."
+    "We are finally here: after having set up the dataset, the model, thee metrics... we are now ready to put all of this together with giotto-deep into a `Trainer` class. Check the cell below."
    ]
   },
   {
@@ -166,18 +168,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# train the model with CV\n",
-    "pipe.train(SGD, 3, cross_validation=True, keep_training=True)\n",
     "\n",
-    "# since we used the keep training flag, the optimiser has not been modified compared to the previous training.\n",
-    "print(pipe.optimizer)"
+    "loss_fn = nn.MSELoss()\n",
+    "\n",
+    "pipe = Trainer(model, (dl_tr, dl_val), loss_fn, writer, l1_norm)\n",
+    "\n",
+    "# train the model with learning rate scheduler\n",
+    "pipe.train(Adam, 3, False, lr_scheduler=ExponentialLR, scheduler_params={\"gamma\": 0.9}, \n",
+    "           profiling=False, store_grad_layer_hist=True, writer_tag=\"line\")\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Extract inner data from your models"
+    "## Checking the results\n",
+    "\n",
+    "Did we managed to get `0.3` after this fisrt part of the training?"
    ]
   },
   {
@@ -186,12 +193,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "pipe.model(torch.tensor([[1,1,1]]).float())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's **train again** the model with cross validation: we just have to set the parameter `cross_validation = True`.\n",
     "\n",
-    "me = ModelExtractor(pipe.model, loss_fn)\n",
-    "\n",
-    "lista = me.get_layers_param()\n",
-    "for k, item in lista.items():\n",
-    "    print(k,item.shape)\n"
+    "The `keep_training = True` flag allow us to resume the trtaining from the same scheduler, optimiser and trained model obtained at the end of the last training in the instance of the class `pipe`."
    ]
   },
   {
@@ -200,16 +211,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "x = next(iter(dl_tr))[0]\n",
-    "list_activations = me.get_activations(x)\n",
-    "len(list_activations)\n"
+    "# train the model with CV\n",
+    "pipe.train(SGD, 3, cross_validation=True, keep_training=True)\n",
+    "\n",
+    "# since we used the keep training flag, the optimiser has not been modified compared to the previous training.\n",
+    "print(pipe.optimizer)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Visualise activations and other topological aspects of your model"
+    "Did we manage this time to get `0.3`?"
    ]
   },
   {
@@ -218,12 +231,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.visualisation import Visualiser\n",
-    "\n",
-    "vs = Visualiser(pipe)\n",
-    "\n",
-    "vs.plot_data_model()\n",
-    "vs.plot_activations(x)\n"
+    "# evaluation\n",
+    "pipe.model(torch.tensor([[1,1,1]]).float())"
    ]
   },
   {
diff --git a/examples/basic_tutorial_tabular.ipynb b/examples/basic_tutorial_tabular.ipynb
index cd432796..9b488268 100644
--- a/examples/basic_tutorial_tabular.ipynb
+++ b/examples/basic_tutorial_tabular.ipynb
@@ -7,7 +7,13 @@
     "# Basic tutorial: tabular data\n",
     "#### Author: Matteo Caorsi\n",
     "\n",
-    "This short tutorial provides you with the basic functioning of *giotto-deep* API.\n",
+    "This short tutorial provides you with the basic functioning of *giotto-deep* API. As an example, we take a simple classification problem in which we have a point cloud representing two entangled tori.\n",
+    "\n",
+    "## Scope\n",
+    "\n",
+    "The goal of this problem is to build a simple feed-forward neural network that is able to separate the two point clouds. The topology of the dataset is non trivial and we empirically discovered that a minimum number of nodes per layer is required to have a good splitting.\n",
+    "\n",
+    "## Plan for this tutorial\n",
     "\n",
     "The main steps of the tutorial are the following:\n",
     " 1. creation of a dataset\n",
@@ -31,29 +37,24 @@
     "import plotly.express as px\n",
     "import torch\n",
     "from torch import nn\n",
+    "from torch.optim.lr_scheduler import ExponentialLR\n",
+    "from torch.optim import SGD, Adam\n",
+    "from torch.utils.data import SubsetRandomSampler\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "from gtda.diagrams import BettiCurve\n",
+    "from gtda.plotting import plot_betti_surfaces\n",
+    "from sklearn.model_selection import StratifiedKFold\n",
     "import pandas as pd\n",
-    "\n",
     "from sklearn import datasets\n",
     "\n",
     "from gdeep.models import FFNet\n",
     "from gdeep.models import ModelExtractor\n",
     "from gdeep.analysis.interpretability import Interpreter\n",
-    "from torch.optim.lr_scheduler import ExponentialLR\n",
-    "from torch.optim import SGD, Adam\n",
     "from gdeep.utility.optimisation import SAMOptimizer\n",
-    "from torch.utils.data import SubsetRandomSampler\n",
-    "\n",
     "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
     "from gdeep.trainer import Trainer\n",
-    "\n",
-    "from torch.utils.tensorboard import SummaryWriter\n",
     "from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder\n",
-    "\n",
-    "from gtda.diagrams import BettiCurve\n",
-    "\n",
-    "from gtda.plotting import plot_betti_surfaces\n",
-    "\n",
-    "from sklearn.model_selection import StratifiedKFold"
+    "from gdeep.visualisation import Visualiser\n"
    ]
   },
   {
@@ -62,14 +63,14 @@
    "source": [
     "# Initialize the tensorboard writer\n",
     "\n",
-    "In order to analyse the reuslts of your models, you need to start tensorboard.\n",
-    "On the terminal, move inside the `/example` folder. There run the following command:\n",
+    "In order to analyse the results of your models, you need to start tensorboard.\n",
+    "On the terminal, move inside the `/exampled` folder. There you can run the following command:\n",
     "\n",
     "```\n",
     "tensorboard --logdir=runs\n",
     "```\n",
     "\n",
-    "Then go [here](http://localhost:6006/) after the training to see all the visualisation results."
+    "Then go [here](http://localhost:6006/) after the training pahse to see all the visualisation results."
    ]
   },
   {
@@ -85,7 +86,10 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Create your dataset"
+    "# Create your dataset\n",
+    "\n",
+    "The dataset consists of two entangled tori and the goal is to classify the points that belong to one or the other torus. Note that the topology of the decision boundary is non trivial!\n",
+    "\n"
    ]
   },
   {
@@ -96,7 +100,7 @@
    "source": [
     "bd = DatasetBuilder(name=\"DoubleTori\")\n",
     "ds_tr, ds_val, _ = bd.build()\n",
-    "#train_indices = list(range(160))\n",
+    "\n",
     "dl_builder = DataLoaderBuilder((ds_tr, ds_val))\n",
     "dl_tr, dl_val, dl_ts = dl_builder.build(({\"batch_size\":23},{\"batch_size\":23}))\n"
    ]
@@ -105,7 +109,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Define and train your model"
+    "## Define and train your model\n",
+    "\n",
+    "In the next few cells we show hoe easy t is in giotto-deep to define a model and start th training: ndeed, once we have the model and the data, only a couple of code lines is needed."
    ]
   },
   {
@@ -135,8 +141,9 @@
    "source": [
     "# initlaise the loss function\n",
     "loss_fn = nn.CrossEntropyLoss()\n",
+    "\n",
     "# initialise the pipelien class\n",
-    "pipe = Trainer(model, (dl_tr, dl_val), loss_fn, writer, StratifiedKFold(3, shuffle=True))\n",
+    "pipe = Trainer(model, (dl_tr, dl_val), loss_fn, writer, k_fold_class=StratifiedKFold(3, shuffle=True))\n",
     "\n",
     "# initialise the SAM optimiser\n",
     "optim = SAMOptimizer(SGD)  # this is a class, not an instance!\n",
@@ -150,11 +157,30 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Add pipeline hooks\n",
+    "## Evaluation\n",
+    "\n",
+    "We can easily evaluate the model performance and compute the confusion matrix:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# evaluation of the model performances for the classification task\n",
+    "pipe.evaluate_classification(2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Advanced topic: add pipeline hooks\n",
     "\n",
     "It is possible to add a hook (a callable) that is called at the end of each training epoch.\n",
     "\n",
-    "The arguments of the hook are fix and have to respect the order!"
+    "The arguments of the hook are fix and have to respect the order you see in this example!"
    ]
   },
   {
@@ -168,7 +194,7 @@
     "    print(f\"We can also get the value of gradients and parameters of the model \"\n",
     "          f\"using the model extractor! {me.get_layers_param():}\")\n",
     "    \n",
-    "\n",
+    "# register the hook\n",
     "pipe.register_pipe_hook(example_of_hook)"
    ]
   },
@@ -198,25 +224,11 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Simply use interpretability tools"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "# Simply use interpretability tools\n",
     "\n",
-    "inter = Interpreter(pipe.model)\n",
-    "inter.interpret_tabular(next(iter(dl_tr))[0], next(iter(dl_tr))[1]);\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Extract inner data from your models"
+    "Eventually, one would like to know why a neural network has made e certain classification choice. Even though the models are complex, there are tools in giotto-deep to help you understand which feature mostly contributed to the network choice.\n",
+    "\n",
+    "In the next celle we use the [Integrated Gradients](https://towardsdatascience.com/integrated-gradients-from-scratch-b46311e4ab4) technique on the model."
    ]
   },
   {
@@ -226,54 +238,22 @@
    "outputs": [],
    "source": [
     "\n",
+    "inter = Interpreter(pipe.model, \"IntegratedGradients\")\n",
+    "inter.interpret_tabular(next(iter(dl_tr))[0], \n",
+    "                        target=next(iter(dl_tr))[1], \n",
+    "                        n_steps=50);\n",
     "\n",
-    "me = ModelExtractor(pipe.model, loss_fn)\n",
-    "\n",
-    "lista = me.get_layers_param()\n",
-    "for k, item in lista.items():\n",
-    "    print(k,item.shape)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "x = next(iter(dl_tr))[0][0]\n",
-    "if x.dtype is not torch.int64:\n",
-    "    res = me.get_decision_boundary(x, n_epochs=100)\n",
-    "    res.shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "x = next(iter(dl_tr))[0]\n",
-    "list_activations = me.get_activations(x)\n",
-    "len(list_activations)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "x, target = next(iter(dl_tr))\n",
-    "if x.dtype is torch.float:\n",
-    "    for gradient in me.get_gradients(x, target=target)[1]:\n",
-    "        print(gradient.shape)"
+    "vs = Visualiser(pipe)\n",
+    "plt = vs.plot_interpreter_tabular(inter)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Visualise activations and other topological aspects of your model"
+    "# Visualise activations and other topological aspects of your model\n",
+    "\n",
+    "This concluding section shows you how it is possible to visualise a lot of nformation on your network in one line thanks to giotto-deep!"
    ]
   },
   {
@@ -282,12 +262,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.visualisation import Visualiser\n",
     "\n",
-    "vs = Visualiser(pipe)\n",
+    "# a batch of data\n",
+    "x = next(iter(dl_tr))[0]\n",
     "\n",
+    "# send to tensorboard the model graph!\n",
     "vs.plot_data_model()\n",
+    "\n",
+    "# send to tensorboard the point cloud, whose coordinates are the activations in each layer, corresponding to the batch x\n",
     "vs.plot_activations(x)\n",
+    "\n",
+    "# plot persistence diagrams of the actiivation point clouds\n",
     "vs.plot_persistence_diagrams(x)\n"
    ]
   },
@@ -297,7 +282,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
+    "# plot the sampled decision boudnary!\n",
     "vs.plot_decision_boundary();"
    ]
   },
@@ -307,6 +292,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# plot the betti curves of each activaton layer\n",
     "vs.betti_plot_layers((0, 1), x)"
    ]
   },
@@ -315,19 +301,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "plt = vs.plot_interpreter_tabular(inter)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# evaluation of the model performances for the classification task\n",
-    "pipe.evaluate_classification(2)"
-   ]
+   "source": []
   },
   {
    "cell_type": "code",
diff --git a/examples/basic_tutorial_textClass.ipynb b/examples/basic_tutorial_text_classification.ipynb
similarity index 68%
rename from examples/basic_tutorial_textClass.ipynb
rename to examples/basic_tutorial_text_classification.ipynb
index 27a3f107..b571ba15 100644
--- a/examples/basic_tutorial_textClass.ipynb
+++ b/examples/basic_tutorial_text_classification.ipynb
@@ -7,10 +7,21 @@
     "# Basic tutorial: text classification\n",
     "#### Author: Matteo Caorsi\n",
     "\n",
-    "This short tutorial provides you with the basic functioning of *giotto-deep* API.\n",
+    "This short tutorial provides you with a basic tutorial on *text classification* using *giotto-deep*.\n",
+    "\n",
+    "## Scope\n",
+    "\n",
+    "Text classification, sometimes called **sentiment analysis** consist is classifying text excerpts into preedefined classes. For example, one may try to classiify movies reviews into five classes with differet level of teh author enjoyment:\n",
+    "\n",
+    "| Movie review | liking score |\n",
+    "|----|-----|\n",
+    "| This movie makes me cry! It's so moving! | 5/5 |\n",
+    "|I hate this movie: bad acting and no plot! Definitely not recommended. | 1/5 |\n",
+    "\n",
+    "## Content\n",
     "\n",
     "The main steps of the tutorial are the following:\n",
-    " 1. creation of a dataset\n",
+    " 1. creation of a dataset and preprocessing\n",
     " 2. creation of a model\n",
     " 3. define metrics and losses\n",
     " 4. run benchmarks\n",
@@ -28,20 +39,27 @@
     "%matplotlib inline\n",
     "\n",
     "import numpy as np\n",
-    "\n",
+    "from gtda.diagrams import BettiCurve\n",
+    "from gtda.plotting import plot_betti_surfaces\n",
     "import torch\n",
     "from torch import nn\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "from torch.utils.data.sampler import SubsetRandomSampler\n",
+    "from torch.optim import Adam, SparseAdam, SGD\n",
+    "# from captum.attr import Occlusion, IntegratedGradients, visualization\n",
+    "from torchvision import transforms\n",
     "\n",
     "from gdeep.models import FFNet\n",
-    "\n",
     "from gdeep.visualisation import  persistence_diagrams_of_activations\n",
-    "\n",
-    "from torch.utils.tensorboard import SummaryWriter\n",
     "from gdeep.data.datasets import DatasetBuilder\n",
-    "\n",
-    "from gtda.diagrams import BettiCurve\n",
-    "\n",
-    "from gtda.plotting import plot_betti_surfaces"
+    "from gdeep.data import PreprocessingPipeline\n",
+    "from gdeep.data import TransformingDataset\n",
+    "from gdeep.data.preprocessors import Normalization, TokenizerTextClassification\n",
+    "from gdeep.data.datasets import DataLoaderBuilder\n",
+    "from gdeep.trainer import Trainer\n",
+    "from gdeep.analysis.interpretability import Interpreter\n",
+    "from gdeep.visualisation import Visualiser\n",
+    "from gdeep.models import ModelExtractor\n"
    ]
   },
   {
@@ -50,8 +68,8 @@
    "source": [
     "# Initialize the tensorboard writer\n",
     "\n",
-    "In order to analyse the reuslts of your models, you need to start tensorboard.\n",
-    "On the terminal, move inside the `/example` folder. There run the following command:\n",
+    "In order to analyse the results of your models, you need to start tensorboard.\n",
+    "On the terminal, move inside the `/examples` folder. There run the following command:\n",
     "\n",
     "```\n",
     "tensorboard --logdir=runs\n",
@@ -73,7 +91,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Create your dataset"
+    "# Create your dataset\n",
+    "\n",
+    "In this section wee will import the [AG_NEWS](https://pytorch.org/text/stable/datasets.html#ag-news) datast. This dataset contains Thomson-Reuters news and they are classified into four macro categories."
    ]
   },
   {
@@ -82,13 +102,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from torch.utils.data.sampler import SubsetRandomSampler\n",
     "\n",
-    "# many time we get an IterableDataset which is good for memory consumption, but cannot be sampled!\n",
-    "# we can load the whole dataset in memory and sample it\n",
+    "# many time we get an IterableDataset which is good for memory consumption, but cannot be subsampled!\n",
+    "# we can entire batches and sample them using the conversion to map type:\n",
     "bd = DatasetBuilder(name=\"AG_NEWS\", convert_to_map_dataset=True)\n",
-    "ds_tr_str, ds_val_str, ds_ts_str = bd.build()\n",
-    "\n"
+    "ds_tr_str, ds_val_str, ds_ts_str = bd.build()\n"
    ]
   },
   {
@@ -107,16 +125,26 @@
     "print(\"Before the preprocessing: \\n\", ds_tr_str[0])"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Required preprocessing\n",
+    "\n",
+    "Neural networks cannot direcly deal with strings. We have first to preprocess the dataset in three main ways:\n",
+    " 1. Tokenise each string into its words (and maybe adjust each word to remove plurals, interjections, capital letters...)\n",
+    " 2. Build a vocabulary out of these tokens (each modified word of point a. is called a token)\n",
+    " 3. Embed each token into a vector, so that each sentence becomes a list of vectors\n",
+    "\n",
+    "The **first two steps** are performed by the `TokenizerTextClassification` class. The embedding will be added directly to the model (see below)."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.data import PreprocessingPipeline\n",
-    "from gdeep.data import TransformingDataset\n",
-    "from gdeep.data.preprocessors import Normalization, TokenizerTextClassification\n",
-    "from gdeep.data.datasets import DataLoaderBuilder\n",
     "\n",
     "ptd = TokenizerTextClassification()\n",
     "\n",
@@ -141,16 +169,18 @@
     "                                                                   \"sampler\":SubsetRandomSampler(train_indices)},\n",
     "                                                                  {\"batch_size\":16,\n",
     "                                                                   \"sampler\":SubsetRandomSampler(test_indices)}\n",
-    "                                                                 ))\n",
-    "\n",
-    "\n"
+    "                                                                 ))\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Define and train your model"
+    "## Define and train your model\n",
+    "\n",
+    "Our model is going to be a simple embedding layer (i.e. a layer whose goal is to transform list of tokens into vectors) , followed by a simple feed-forward layer.\n",
+    "\n",
+    "Furtheremore, snce the output of te embedding layer has three dimensions, we need to remove one before the tensor can be fed to the final feed-forward layer: this is - in this very smple tutoral - via an average operation."
    ]
   },
   {
@@ -159,11 +189,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.trainer import Trainer\n",
-    "from torch.optim import Adam, SparseAdam, SGD\n",
-    "from captum.attr import Occlusion, IntegratedGradients, visualization\n",
-    "from torchvision import transforms\n",
-    "\n",
     "\n",
     "class TextClassificationModel(nn.Module):\n",
     "\n",
@@ -185,6 +210,13 @@
     "        return self.fc(mean)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In the next cell we initialise the model parameters and the model as well."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -199,26 +231,14 @@
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [],
    "source": [
+    "# Training the model\n",
     "\n",
-    "loss_fn = nn.CrossEntropyLoss()\n",
-    "\n",
-    "pipe = Trainer(model, (dl_tr2, dl_ts2), loss_fn, writer)\n",
+    "We are finally ready to train our model: we have already built a dataset (that is compatble wth thhe model). We also know that the task is a classification task, hence we can use the CrossEntropy loss!\n",
     "\n",
-    "# train the model\n",
-    "pipe.train(SGD, 7, False, {\"lr\":0.01}, {\"batch_size\":20})\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Simply use interpretability tools"
+    "Let's see how in giotto-deep, in only a couple of lines, one can run the training."
    ]
   },
   {
@@ -227,61 +247,25 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from gdeep.analysis.interpretability import Interpreter\n",
-    "from gdeep.visualisation import Visualiser\n",
     "\n",
-    "vs = Visualiser(pipe)\n",
-    "inter = Interpreter(pipe.model, method=\"LayerIntegratedGradients\")\n",
+    "loss_fn = nn.CrossEntropyLoss()\n",
     "\n",
-    "inter.interpret_text(\"I am writing about money and business\", 0, ptd.vocabulary, ptd.tokenizer, layer=pipe.model.embedding)\n",
+    "pipe = Trainer(model, (dl_tr2, dl_ts2), loss_fn, writer)\n",
     "\n",
-    "vs.plot_interpreter_text(inter);"
+    "# train the model\n",
+    "pipe.train(SGD, 7, False, {\"lr\":0.01}, {\"batch_size\":20})\n",
+    "\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Extract inner data from your models"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from gdeep.models import ModelExtractor\n",
+    "# Simply use interpretability tools\n",
     "\n",
-    "me = ModelExtractor(model, loss_fn)\n",
+    "In the next two cells we use interpretability tool on the model to see what and how our model has learned. Basically, we want to understand how each word in a seentence contributes to the classification choice.\n",
     "\n",
-    "lista = me.get_layers_param()\n",
-    "\n",
-    "for k, item in lista.items():\n",
-    "    print(k,item.shape)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "x = next(iter(dl_tr2))[0][0]\n",
-    "if x.dtype is not torch.int64:\n",
-    "    res = me.get_decision_boundary(x, n_epochs=1)\n",
-    "    res.shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "x = next(iter(dl_tr2))[0]\n",
-    "list_activations = me.get_activations(x)\n",
-    "len(list_activations)\n"
+    "But frst, how about sending to the tensorboard the model structure an interactiive inspecton? It's as simple as writing the code in th next cell."
    ]
   },
   {
@@ -290,17 +274,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "x, target = next(iter(dl_tr2))\n",
-    "if x.dtype is torch.float:\n",
-    "    for gradient in me.get_gradients(x, target=target)[1]:\n",
-    "        print(gradient.shape)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Visualise activations and other topological aspects of your model"
+    "vs = Visualiser(pipe)\n",
+    "vs.plot_data_model()  "
    ]
   },
   {
@@ -309,19 +284,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Interporetability!\n",
     "\n",
-    "vs.plot_data_model()\n",
-    "vs.plot_activations(x)\n",
-    "vs.plot_persistence_diagrams(x)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "vs.betti_plot_layers((0, 1), x)"
+    "inter = Interpreter(pipe.model, method=\"LayerIntegratedGradients\")\n",
+    "\n",
+    "inter.interpret_text(\"I am writing about money and business\", 0, ptd.vocabulary, ptd.tokenizer, layer=pipe.model.embedding)\n",
+    "\n",
+    "\n",
+    "vs.plot_interpreter_text(inter);"
    ]
   },
   {
diff --git a/examples/basic_tutorial_translation.ipynb b/examples/basic_tutorial_translation.ipynb
index 9d10eca7..97ece72c 100644
--- a/examples/basic_tutorial_translation.ipynb
+++ b/examples/basic_tutorial_translation.ipynb
@@ -486,96 +486,6 @@
     "# go to TB to see the results + screenshot"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Appendix: Extract inner data from your models\n",
-    "\n",
-    "With `gdeep` is is pretty straight forwasrd to extract data about the inner working of the models.\n",
-    "\n",
-    "For example you can use the `ModelExtract` and get:\n",
-    " - the parameters of each layer\n",
-    " - the activation functions (given an input)\n",
-    " - the gradients, given a target"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\n",
-    "# initialise the ModelExtractor\n",
-    "me = ModelExtractor(pipe.model, loss_fn)\n",
-    "\n",
-    "lista = me.get_layers_param()\n",
-    "\n",
-    "for k, item in lista.items():\n",
-    "    print(k,item.shape)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "DEVICE = torch.device(\"cpu\")\n",
-    "x = next(iter(dl_tr))[0]\n",
-    "pipe.model.eval()\n",
-    "pipe.model(x.to(DEVICE))\n",
-    "\n",
-    "list_activations = me.get_activations(x)\n",
-    "len(list_activations)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "x = next(iter(dl_tr))[0][0]\n",
-    "if x.dtype is not torch.int64:\n",
-    "    res = me.get_decision_boundary(x, n_epochs=1)\n",
-    "    res.shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "x, target = next(iter(dl_tr))\n",
-    "if x.dtype is torch.float:\n",
-    "    for gradient in me.get_gradients(x, target=target)[1]:\n",
-    "        print(gradient.shape)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "You can also visualise the model graph on tensorboard interactively!"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\n",
-    "# initialise the Visualiser\n",
-    "vs = Visualiser(pipe)\n",
-    "\n",
-    "vs.plot_data_model()\n",
-    "\n"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
diff --git a/examples/images/image_class.png b/examples/images/image_class.png
new file mode 100644
index 00000000..10e7132d
Binary files /dev/null and b/examples/images/image_class.png differ
diff --git a/examples/images/regression.jpeg b/examples/images/regression.jpeg
new file mode 100644
index 00000000..099dfba0
Binary files /dev/null and b/examples/images/regression.jpeg differ
diff --git a/examples/images/translation_transformer.png b/examples/images/translation_transformer.png
new file mode 100644
index 00000000..0529c0fe
Binary files /dev/null and b/examples/images/translation_transformer.png differ
diff --git a/examples/optimising_topology.ipynb b/examples/optimising_topology.ipynb
index 9628e6cc..3f70e7e9 100644
--- a/examples/optimising_topology.ipynb
+++ b/examples/optimising_topology.ipynb
@@ -143,7 +143,7 @@
     "pg = PersistenceGradient(homology_dimensions=hom_dim,\n",
     "                         zeta=0.1, max_edge_length=0.5,\n",
     "                         collapse_edges=False)\n",
-    "fig, fig3d, loss_val = pg.SGD(Xp,0.07,32)\n",
+    "fig, fig3d, loss_val = pg.sgd(Xp,0.07,32)\n",
     "\n",
     "# plot of the evolution over the GD iterations\n",
     "fig.show()\n",
@@ -235,7 +235,7 @@
     "pg = PersistenceGradient(homology_dimensions=hom_dim,\n",
     "                         zeta=0.0,collapse_edges=False,\n",
     "                         metric=\"precomputed\")\n",
-    "fig, fig3d, loss_val = pg.SGD(dist,0.1,10)\n",
+    "fig, fig3d, loss_val = pg.sgd(dist,0.1,10)\n",
     "# plot of the loss function\n",
     "px.line(loss_val)"
    ]
@@ -286,7 +286,7 @@
     "pg = PersistenceGradient(homology_dimensions=hom_dim,\n",
     "                         zeta=0.01,collapse_edges=False,\n",
     "                         metric=\"precomputed\", directed = True)\n",
-    "fig, fig3d, loss_val = pg.SGD(dist,0.1,10)\n",
+    "fig, fig3d, loss_val = pg.sgd(dist,0.1,10)\n",
     "# plot of the loss function\n",
     "px.line(loss_val)"
    ]
diff --git a/gdeep/analysis/interpretability/__init__.py b/gdeep/analysis/interpretability/__init__.py
index 20e1680c..69a4231a 100644
--- a/gdeep/analysis/interpretability/__init__.py
+++ b/gdeep/analysis/interpretability/__init__.py
@@ -1,6 +1,9 @@
 from .captum import Interpreter
+from .attribution_factory import AttributionFactory, get_attr
 
 
 __all__ = [
-    'Interpreter'
+    'Interpreter',
+    'AttributionFactory',
+    'get_attr'
     ]
diff --git a/gdeep/analysis/interpretability/attribution_factory.py b/gdeep/analysis/interpretability/attribution_factory.py
new file mode 100644
index 00000000..f2c62fe0
--- /dev/null
+++ b/gdeep/analysis/interpretability/attribution_factory.py
@@ -0,0 +1,73 @@
+
+from typing import Dict, Any
+
+from captum import attr
+
+# Define the dataset factory class for the tori dataset and torchvision datasets
+# using the factory design pattern
+# https://realpython.com/factory-method-python/
+
+
+class AttributionFactory(object):
+    """ Attribution factory class for the Captum integration.
+    This factory will contain the attributions techniques of captum.
+    """
+
+    _builders: Dict[str, Any] = {}
+
+    def __init__(self):
+        pass
+        
+    def register_builder(self, key: str, builder: Any):
+        """this method adds to the internal builders dictionary
+        new dataloader builders
+        """
+        self._builders[key] = builder    # type: ignore
+        
+    def build(self, key: str, *args, **kwargs) -> Any:  # type: ignore
+        """This method returns the DataLoader builder
+        corresponding to the input key.
+
+        Args:
+            key:
+                the name of the dataset
+        """
+        builder = self._builders.get(key)    # type: ignore
+        if builder is None:
+            raise ValueError(f"No builder registered for key {key}")
+        return builder(*args, **kwargs)  # type: ignore
+
+
+class AttributionBuilder(object):
+    """Builder class for the torchvision dataset
+        and all its variations"""
+
+    def __init__(self, attr_name: str):
+        self.attr_name = attr_name
+
+    def __call__(self, *args, **kwargs):  # type: ignore
+        return attr.__getattribute__(self.attr_name)(  # type: ignore
+            *args, **kwargs  # type: ignore
+        )
+
+
+def get_attr(key: str, *args, **kwargs) -> Any:  # type: ignore
+    """Get a dataset from the factory
+
+    Args:
+        key :
+            The name of the dataset, corresponding
+            to the key in the list of builders
+        **kwargs:
+            The keyword arguments to pass to the dataset builder
+
+    Returns:
+        captum.attr:
+            The captum interpretability tool
+    """
+    factory = AttributionFactory()
+
+    for attr_name in attr.__all__:
+        factory.register_builder(attr_name, AttributionBuilder(attr_name))
+
+    return factory.build(key, *args, **kwargs)  # type: ignore
diff --git a/gdeep/analysis/interpretability/captum.py b/gdeep/analysis/interpretability/captum.py
index c46c9480..3283d8b7 100644
--- a/gdeep/analysis/interpretability/captum.py
+++ b/gdeep/analysis/interpretability/captum.py
@@ -1,49 +1,53 @@
-from captum.attr import * # TokenReferenceBase, \
-    # visualization, FeatureAblation, \
-    # DeepLift, NoiseTunnel, IntegratedGradients, \
-    # LayerIntegratedGradients, Occlusion, \
-    # GuidedGradCam, Saliency, InputXGradient
+from typing import Any, Tuple, Dict, Callable, List, Optional
+
+from captum.attr import *
 import torch
+from torch import nn
+from .attribution_factory import get_attr
+
+from gdeep.utility import DEVICE
+
+Tensor = torch.Tensor
 
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-else:
-    DEVICE = torch.device("cpu")
 
 class Interpreter:
     """Class to visualise the activation maps,
-    the attribution maps and salicy maps using
+    the attribution maps and saliency maps using
     different techniques.
 
     Args:
-        model (nn.Module):
+        model:
             the standard pytorch model
-        method (string):
+        method:
             the interpretability method. Find
             more info at https://captum.ai/tutorials/
 
     """
 
-    def __init__(self, model,
-                 method="IntegratedGradients"):
+    def __init__(self, model: nn.Module,
+                 method: str = "IntegratedGradients"):
         # self.model = model
         self.model = model.to(DEVICE)
         self.method = method
         self.stored_visualisations = []
         self.image = None
-        self.X = None
+        self.x = None
         self.sentence = None
-        self.attrib = None
+        self.attribution = None
 
-    def interpret_image(self, X, y, layer = None, **kwargs):
+    def interpret_image(self,
+                        x: Tensor,
+                        y: Any,
+                        layer: Optional[torch.nn.Module] = None,
+                        **kwargs: Any) -> Tuple[Tensor, Tensor]:
         """This method creates an image interpreter. This class
         is based on captum.
 
         Args:
-            X (tensor):
+            x:
                 the tensor corresponding to the input image.
                 It is expected to be of size ``(b, c, h, w)``
-            y (int or float or str)
+            y:
                 the label we want to check the interpretability
                 of.
             layer (nn.Module, optional):
@@ -51,102 +55,76 @@ def interpret_image(self, X, y, layer = None, **kwargs):
                 of self.model
 
         Returns
-            (torch.Tensor, torch.Tensor):
+            torch.Tensor, torch.Tensor:
                 the input image and the attribution
                 image respectively.
         """
-        self.X = X.to(DEVICE)
-        if self.method in ("GuidedGradCam",
-                           "LayerConductance",
-                           "LayerActivation",
-                           "LayerGradCam",
-                           "LayerDeepLift",
-                           "LayerFeatureAblation",
-                           "LayerIntegratedGradients",
-                           "LayerGradientShap",
-                           "LayerDeepLiftShap"):
-            occlusion = eval(self.method+"(self.model, layer)")
-        else:
-            occlusion = eval(self.method+"(self.model)")
+        self.x = x.to(DEVICE)
+        attr_class = get_attr(self.method, self.model, layer)
         self.model.eval()
-        att = occlusion.attribute(self.X, target=y, **kwargs)
-        self.image = self.X
-        self.attrib = att
-        return self.X, att
+        self.attribution = attr_class.attribute(self.x, target=y, **kwargs)
+        self.image = self.x
+        return self.x, self.attribution
 
-    def interpret_tabular(self, X_test, y, **kwargs):
-        """This method creates an image interpreter. This class
+    def interpret_tabular(self, x: Tensor, **kwargs: Any) -> None:
+        """This method creates a tabular interpreter. This class
         is based on captum.
 
         Args:
-            X (tensor):
+            x :
                 the tensor corresponding to the input image.
                 It is expected to be of size ``(b, c, h, w)``
-            y (int or float or str)
-                the label we want to check the interpretability
-                of.
 
-        Returns
-            (torch.Tensor, torch.Tensor):
-                the input image and the attribution
-                image respectively.
         """
-        self.X = X_test.to(DEVICE)  # needed for plotting functions
-        y = y.to(DEVICE)
+        self.x = x.to(DEVICE)  # needed for plotting functions
         self.model.eval()
-        ig = IntegratedGradients(self.model)
-        ig_nt = NoiseTunnel(ig)
-        dl = DeepLift(self.model)
+        attr_class = get_attr(self.method, self.model)
+        self.attribution = attr_class.attribute(self.x,
+                                                **kwargs)
+        # ig = IntegratedGradients(self.model)
+        # ig_nt = NoiseTunnel(ig)
+        # dl = DeepLift(self.model)
         # gs = GradientShap(self.model)
-        fa = FeatureAblation(self.model)
-        self.ig_attr_test = ig.attribute(self.X,
-                                         n_steps=50,
-                                         target=y,
-                                         **kwargs)
-        self.ig_nt_attr_test = ig_nt.attribute(self.X,
-                                               target=y,
-                                               **kwargs)
-        self.dl_attr_test = dl.attribute(self.X,
-                                         target=y,
-                                         **kwargs)
-        # self.gs_attr_test = gs.attribute(X_test, X_train, **kwargs)
-        self.fa_attr_test = fa.attribute(self.X,
-                                         **kwargs)
-
-    def interpret_text(self, sentence, label, vocab,
-                       tokenizer, layer,
-                       min_len=7):
-        """This method creates an image interpreter. This class
+        # fa = FeatureAblation(self.model)
+        # self.ig_attr_test = ig.attribute(self.x,
+        #                                  n_steps=50,
+        #                                  target=y,
+        #                                  **kwargs)
+
+    def interpret_text(self, sentence: str,
+                       label: Any,
+                       vocab: Dict[str, int],
+                       tokenizer: Callable[[str], List[str]],
+                       layer: torch.nn.Module,
+                       min_len: int = 7) -> None:
+        """This method creates a text interpreter. This class
         is based on captum.
 
         Args:
-            sentence (string):
+            sentence :
                 the input sentence
-            label (int or float or str)
+            label:
                 the label we want to check the interpretability
                 of.
-            vocab (vocabulary):
-                a ``gdeep.data.PreprocessText`` vocabulary. Can
-                be extracted via the ``vocabulary``attribute.
-            tokenizer (tokenizer):
-                a ``gdeep.data.PreprocessText`` tokenizer. Can
-                be extracted via the ``tokenizer``attribute.
-            layer (nn.Module):
-                torch module correspondign to the layer belonging to
+            vocab :
+                a ``gdeep.data.preprocessors`` vocabulary. Can
+                be extracted via the ``vocabulary`` attribute.
+            tokenizer :
+                a ``gdeep.data.preprocessors`` tokenizer. Can
+                be extracted via the ``tokenizer`` attribute.
+            layer :
+                torch module corresponding to the layer belonging to
                 ``self.model``.
+            min_len:
+                minimum length of the text. Shorter texts are padded
 
-        Returns
-            (torch.Tensor, torch.Tensor):
-                the input image and the attribution
-                image respectively.
         """
         self.model.eval()
         self.sentence = sentence
-        lig = eval(self.method+"(" + ",".join(("self.model", #"self.model." +
-                   "layer")) + ")")
+        attr_class = get_attr(self.method, self.model, layer)
         text = tokenizer(sentence)
         if len(text) < min_len:
-            text += ['.'] * (min_len - len(text))
+            text += [' '] * (min_len - len(text))
         indexed = [vocab[t] for t in text]
 
         self.model.zero_grad()
@@ -162,31 +140,34 @@ def interpret_text(self, sentence, label, vocab,
         pred = torch.max(pred_temp)
         pred_ind = torch.argmax(pred_temp).item()
         # generate reference indices for each sample
-        PAD_IND = vocab['.']
-        token_reference = TokenReferenceBase(reference_token_idx=PAD_IND)
+        pad_index = 0
+        token_reference = TokenReferenceBase(reference_token_idx=pad_index)
         reference_indices = \
             token_reference.generate_reference(seq_length,
                                                DEVICE).unsqueeze(0)
         # compute attributions and approximation
         # delta using layer integrated gradients
-        attributions_ig, delta = lig.attribute(input_indices,
-                                               reference_indices,
-                                               target=label,
-                                               n_steps=500,
-                                               return_convergence_delta=True)
+        self.attribution, delta = attr_class.attribute(input_indices,
+                                                       reference_indices,
+                                                       target=label,
+                                                       n_steps=500,
+                                                       return_convergence_delta=True)
 
         print('pred: ', pred_ind, '(', '%.2f' % pred, ')',
               ', delta: ', abs(delta.item()))
 
-        self.add_attributions_to_visualizer(attributions_ig, text, pred,
+        self.add_attributions_to_visualizer(self.attribution, text, pred,
                                             pred_ind, label, delta.item(),
-                                            self.stored_visualisations,
-                                            vocab)
+                                            self.stored_visualisations)
 
     @staticmethod
-    def add_attributions_to_visualizer(attributions, text, pred, pred_ind,
-                                       label, delta, vis_data_records,
-                                       vocab):
+    def add_attributions_to_visualizer(attributions,
+                                       text: List[str],
+                                       pred: Tensor,
+                                       pred_ind: Tensor,
+                                       label: Any,
+                                       delta: Any,
+                                       vis_data_records: List) -> None:
         attributions = attributions.sum(dim=2).squeeze(0)
         attributions = attributions / torch.norm(attributions)
         attributions = attributions.cpu().detach().numpy()
diff --git a/gdeep/models/__init__.py b/gdeep/models/__init__.py
index 6e52cf54..83227176 100644
--- a/gdeep/models/__init__.py
+++ b/gdeep/models/__init__.py
@@ -1,18 +1,15 @@
 
 from .simple_nn import FFNet
-from .utility import SaveOutput, SaveNodeOutput, \
-    SaveLayerOutput, PeriodicNeuralNetwork
+from .utils import SaveOutput, SaveNodeOutput, \
+    SaveLayerOutput
+from .periodic_nn import PeriodicNeuralNetwork
 from .extractor import ModelExtractor
 
 __all__ = [
     'FFNet',
-    'CircleNN',
-    'CircleNN3D',
     'SaveOutput',
     'SaveNodeOutput',
     'SaveLayerOutput',
-    'get_activations',
-    'LayersList',
     'PeriodicNeuralNetwork',
     'ModelExtractor'
     ]
diff --git a/gdeep/models/extractor.py b/gdeep/models/extractor.py
index 4e1a34ab..f61e65fc 100644
--- a/gdeep/models/extractor.py
+++ b/gdeep/models/extractor.py
@@ -1,32 +1,35 @@
+from typing import List, Any, Dict, Callable, Tuple
+
 import torch
 
 from ..analysis.decision_boundary import GradientFlowDecisionBoundaryCalculator
 from ..analysis.decision_boundary import UniformlySampledPoint
 from . import SaveLayerOutput
+from gdeep.utility import DEVICE
+
+Tensor = torch.Tensor
 
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-else:
-    DEVICE = torch.device("cpu")
 
 class ModelExtractor:
     """This class wraps nn.Modules to extract
     weights, activations, gradients and decision boundaries
 
     Args:
-        model (nn.Module):
+        model:
             standard torch module
-        loss_fn (Callable):
+        loss_fn :
             loss function
     """
 
-    def __init__(self, model, loss_fn):
+    def __init__(self, model: torch.nn.Module,
+                 loss_fn: Callable[[Tensor, Tensor], Tensor]) -> None:
         # self.model = model
         self.model = model.to(DEVICE)
         self.loss_fn = loss_fn
 
-    def get_decision_boundary(self, input_example, n_epochs=100,
-                              precision=0.1):
+    def get_decision_boundary(self, input_example: Tensor,
+                              n_epochs: int=100,
+                              precision: float=0.1) -> Tensor:
         """Compute the decision boundary for self.model
         with self.loss_fn
         
@@ -66,7 +69,7 @@ def get_decision_boundary(self, input_example, n_epochs=100,
                                                 precision).detach()
         return res
 
-    def get_activations(self, X):
+    def get_activations(self, x: Tensor) -> List[Tensor]:
         """Compute the activations of self.model with input
         `X`
 
@@ -89,7 +92,7 @@ def get_activations(self, X):
             hook_handles.append(handle)
 
         self.model.eval()
-        self.model(X.to(DEVICE))
+        self.model(x.to(DEVICE))
 
 
         for handle in hook_handles:
@@ -98,7 +101,7 @@ def get_activations(self, X):
         self.model.train()
         return saved_output_layers.get_outputs()
 
-    def get_layers_param(self) -> dict:
+    def get_layers_param(self) -> Dict[str, Tensor]:
         """Returns parameters of layers
 
         Returns:
@@ -113,7 +116,7 @@ def get_layers_param(self) -> dict:
         #    layer_data[name]=layer_param.data
         return self.model.state_dict()
         
-    def get_layers_grads(self) -> dict:
+    def get_layers_grads(self) -> List[Tensor]:
         """Returns the gradients of each layer
 
         Returns:
@@ -131,20 +134,19 @@ def get_layers_grads(self) -> dict:
             output.append(v.grad)
         return output
 
-    def get_gradients(self, x, **kwargs) -> tuple:
+    def get_gradients(self, x: Tensor, **kwargs) \
+            -> Tuple[Tensor, List[Tensor]]:
         """Returns the **averaged gradient** of the self.loss_fn.
         To specify the target variable (e.g. the true class),
         use the keywords argument `target=`
 
         Args:
-            x (Tensor):
+            x :
                 point at which to compute grad
 
         Returns:
-            tensor:
-                the gradients
-            list:
-                list of tensors, corresponding
+            tensor,  list:
+                the gradients; the list of tensors, corresponding
                 to the gradient of the weights.
         """
         
diff --git a/gdeep/models/periodic_nn.py b/gdeep/models/periodic_nn.py
new file mode 100644
index 00000000..e1fa24e4
--- /dev/null
+++ b/gdeep/models/periodic_nn.py
@@ -0,0 +1,43 @@
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+Tensor = torch.Tensor
+
+class PeriodicNeuralNetwork(nn.Module):
+    """Makes a periodic `nn.Module` of `nn`. `boundary_list` specifies
+    the elementary region the periodic neural network is supported on.
+    This creates neural networks
+    factorized by the orbits of a group effect,
+    which is given by reflection at the
+    edges of the elementary domain.
+    This class can be interpreted as an adapter for periodic neural networks.
+
+    Args:
+        nn :
+            the standard torch module, your network
+        boundary_list :
+            list of pairs of floats, each
+            defining the boundaries of a hypercube
+    """
+
+    def __init__(self, nn: nn.Module,
+                 boundary_list:List[Tuple[float, float]]) -> None:
+        super().__init__()
+        self.nn = nn
+        self.interval_length = torch.tensor([[b-a for a, b
+                                              in boundary_list]])
+        self.left_interval_bound = torch.tensor([[a for a, b
+                                                  in boundary_list]])
+
+    def forward(self, x_cont:Tensor) -> Tensor:
+        x_cont = torch.abs(
+                    torch.remainder(
+                        x_cont-self.left_interval_bound,
+                        2.*self.interval_length
+                        ) - self.interval_length
+                        ) + self.left_interval_bound
+        x_cont = self.nn.forward(x_cont)
+        return x_cont
\ No newline at end of file
diff --git a/gdeep/models/simple_nn.py b/gdeep/models/simple_nn.py
index 7f7aa8ca..47819438 100644
--- a/gdeep/models/simple_nn.py
+++ b/gdeep/models/simple_nn.py
@@ -27,45 +27,3 @@ def forward(self, x):
         x = self.linears[-1](x)
         return x
 
-#    def __init__(self, verbose=0, arch=[2, 3, 3, 2]):
-#        super(FFNet, self).__init__()
-#        self.verbose = verbose
-#        self.arch = arch
-#        for i, in_f in enumerate(arch):
-#            if i < len(arch) - 1:
-#                val = "self.layer" + str(i) + "=" + \
-#                   "nn.Linear("+str(in_f) +","+str(arch[i+1])+")"
-#                exec(val)
-#                val2 = "self.layer"+str(i)+".weight.data.uniform_(-1, 1)"
-#                eval(val2)
-#                val3 = "self.layer"+str(i)+".bias.data.uniform_(-1, 1)"
-#                eval(val3)
-
-#    def forward(self, x_cont):
-#        # output_vars = []
-#        for i, in_f in enumerate(self.arch):
-#            if i < len(self.arch) - 1:
-#                if i == 0:
-#                    val = "x" + str(i) + "=F.relu(" + \
-#                        "self.layer" + str(i) + "(x_cont)" + ")"
-#                # put softmax in last layer
-#                elif i == len(self.arch) - 2:
-#                    val = "x"+str(i)+"=F.softmax(" + \
-#                        "self.layer"+str(i)+"(x" + str(i - 1) + \
-#                            ")" + ",dim=-1)"
-#                else:
-#                    val = "x" + str(i) + "=F.relu(" + \
-#                        "self.layer" + str(i) + "(x" + \
-#                            str(i - 1) + ")" + ")"
-#                if self.verbose:
-#                    print(val)
-#                try:
-#                    exec(val)
-#                except:
-#                    raise ValueError("The dimensions of x" + str(i)
-#                                     + " are not correct!")
-#                if self.verbose:
-#                    print("x" + str(i) + " size: ",
-#                          eval("x" + str(i) + ".shape"))
-#        i -= 1
-#        return eval("x" + str(i))
diff --git a/gdeep/models/utility.py b/gdeep/models/utility.py
deleted file mode 100644
index 8aaa0f07..00000000
--- a/gdeep/models/utility.py
+++ /dev/null
@@ -1,85 +0,0 @@
-import torch
-from torch import nn
-
-
-class SaveOutput:
-    """ General class for saving outputs
-    outputs will be stored in a list 'outputs'
-    """
-    def __init__(self):
-        self.outputs = []
-
-    def __call__(self, module, module_in, module_out):
-        pass
-
-    def clear(self):
-        self.outputs = []
-
-    def get_outputs(self):
-        return self.outputs
-
-
-class SaveNodeOutput(SaveOutput):
-    """ Class for saving activations of a node in
-    a neural network.
-    """
-    def __init__(self, entry=0):
-        super().__init__()
-        self.entry = entry
-
-    def __call__(self, module, module_in, module_out):
-        self.outputs.append(module_out[:, self.entry].detach())
-
-
-class SaveLayerOutput(SaveOutput):
-    def __call__(self, module: nn.Module, module_in, module_out):
-        """ Add activation 
-
-        Args:
-            module ([type]):
-                [description]
-            module_in ([type]):
-                [description]
-            module_out ([type]):
-                [description]
-        """
-        if isinstance(module_out, tuple):
-            self.outputs.append(module_out[0].detach())
-        else:
-            self.outputs.append(module_out.detach())
-
-
-class PeriodicNeuralNetwork(nn.Module):
-    """Makes a periodic `nn.Module` of `nn`. `boundary_list` specifies
-    the elementary region the periodic neural network is supported on.
-    This creates neural networks
-    factorized by the orbits of a group effect,
-    which is given by reflection at the
-    edges of the elementary domain.
-    This class can be interpreted as an adapter for periodic neural networks.
-
-    Args:
-        nn (nn.Module):
-            the standard torch module, your network
-        boundary_list (list):
-            list of pairs of floats, each
-            defining the boundaries of a hypercube
-    """
-
-    def __init__(self, nn: nn.Module, boundary_list):
-        super().__init__()
-        self.nn = nn
-        self.interval_length = torch.tensor([[b-a for a, b
-                                              in boundary_list]])
-        self.left_interval_bound = torch.tensor([[a for a, b
-                                                  in boundary_list]])
-
-    def forward(self, x_cont):
-        x_cont = torch.abs(
-                    torch.remainder(
-                        x_cont-self.left_interval_bound,
-                        2.*self.interval_length
-                        ) - self.interval_length
-                        ) + self.left_interval_bound
-        x_cont = self.nn.forward(x_cont)
-        return x_cont
diff --git a/gdeep/models/utils.py b/gdeep/models/utils.py
new file mode 100644
index 00000000..6a4e7b9c
--- /dev/null
+++ b/gdeep/models/utils.py
@@ -0,0 +1,65 @@
+from abc import ABC, abstractmethod
+from typing import List, Any, Union
+
+import torch
+from torch import nn
+
+Tensor = torch.Tensor
+
+
+class SaveOutput(ABC):
+    """ General class for saving outputs of
+    neural networks.
+    Outputs will be stored in a list 'outputs'
+    """
+    outputs: List[Any]
+
+    @abstractmethod
+    def __call__(self, module: nn.Module, module_in: Tensor, module_out: Union[List[Tensor], Tensor]) -> None:
+        pass
+
+    def clear(self) -> None:
+        self.outputs = []
+
+    def get_outputs(self) -> List[Any]:
+        return self.outputs
+
+
+class SaveNodeOutput(SaveOutput):
+    """ Class for saving activations of a node in
+    a neural network.
+
+    Args:
+        entry:
+            select the entry row in the output
+            of the layer, i.e. the node output
+    """
+
+    def __init__(self, entry: int = 0):
+        self.outputs = []
+        self.entry = entry
+
+    def __call__(self, module, module_in: Tensor, module_out: Tensor):
+        self.outputs.append(module_out[:, self.entry].detach())
+
+
+class SaveLayerOutput(SaveOutput):
+    def __init__(self) -> None:
+        self.outputs = []
+
+    def __call__(self, module: nn.Module, module_in: Tensor,
+                 module_out: Union[List[Tensor], Tensor]) -> None:
+        """This function stores the layer activations
+
+        Args:
+            module :
+                the neural network model (nn.Module)
+            module_in:
+                the input tensor of the module
+            module_out:
+                the output tensor of the module
+        """
+        if isinstance(module_out, tuple):
+            self.outputs.append(module_out[0].detach())
+        else:
+            self.outputs.append(module_out.detach())
diff --git a/gdeep/search/__init__.py b/gdeep/search/__init__.py
index 05de8526..8b8a5d59 100644
--- a/gdeep/search/__init__.py
+++ b/gdeep/search/__init__.py
@@ -1,11 +1,14 @@
 
-from .benchmark import Benchmark, _benchmarking_param
+from .benchmark import Benchmark, _benchmarking_param, TrainerConfig
 from .hpo import HyperParameterOptimization, GiottoSummaryWriter
 from .persformer_hyperparamter_search import PersformerHyperparameterSearch
 from ._utils import clean_up_files
+from .hpo_config import HPOConfig
 
 __all__ = [
     'Benchmark',
+    'TrainerConfig',
+    'HPOConfig',
     'HyperParameterOptimization',
     'GiottoSummaryWriter',
     '_benchmarking_param',
diff --git a/gdeep/search/benchmark.py b/gdeep/search/benchmark.py
index f3fa6be4..9f81c012 100644
--- a/gdeep/search/benchmark.py
+++ b/gdeep/search/benchmark.py
@@ -1,40 +1,46 @@
 import warnings
 from typing import Tuple, Any, Dict, Callable, \
-    Type, Optional, List
+    Type, Optional, List, Union
 
 import torch
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
-from sklearn.model_selection._split import BaseCrossValidator
+from sklearn.model_selection._split import BaseCrossValidator  # noqa
 from torch.optim import Optimizer
-from torch.optim.lr_scheduler import _LRScheduler
+from torch.optim.lr_scheduler import _LRScheduler  # noqa
+from optuna.trial._base import BaseTrial  # noqa
+from sklearn.model_selection import KFold
 
 from gdeep.trainer import Trainer
-from gdeep.utility import _are_compatible
-from sklearn.model_selection import KFold
-from gdeep.utility import DEVICE
+from gdeep.utility import _are_compatible  # noqa
+from gdeep.trainer import accuracy, TrainerConfig
 
 Tensor = torch.Tensor
 
+
 class Benchmark:
     """This is the generic class that allows
     the user to perform benchmarking over different
     datasets and models.
 
     Args:
-        models_dicts (list of dicts):
-            each dictionary has two items, `"name":"string"`
+        models_dicts :
+            each dictionary has two items, ``"name":"string"``
             the name of the model and `"model":nn.Module` a
             standard torch model
-        dataloaders_dicts (utils.DataLoader):
-            each dictionary has two items, `"name":"string"`
-            the name of the model and `"dataloaders":list` a
-            list of standard torch.dataloaders, e.g. `(dl_tr, dl_ts)`
-        loss_fn (Callables):
+        dataloaders_dicts :
+            each dictionary has two items, ``"name":"string"``
+            the name of the model and ``"dataloaders":list`` a
+            list of standard torch.dataloaders, e.g. ``(dl_tr, dl_ts)``
+        loss_fn:
             loss function
-        writer (tensorboard SummaryWriter):
+        writer:
             tensorboard writer
-        KFold_class (sklearn.model_selection, default KFold()):
+        training_metric:
+            the function that computes the metric: it shall
+            have two arguments, one for the prediction
+            and the other for the ground truth
+        k_fold_class:
             the class instance to implement the KFold, can be
             any of the Splitter classes of sklearn. More
             info at https://scikit-learn.org/stable/modules/classes.html#module-sklearn.model_selection
@@ -42,23 +48,27 @@ class Benchmark:
     """
 
     def __init__(self,
-                 models_dicts:Dict[str, torch.nn.Module],
-                 dataloaders_dicts:Dict[str, List[DataLoader[Tuple[Tensor, Tensor]]]],
+                 models_dicts: Dict[str, Union[torch.nn.Module, str]],
+                 dataloaders_dicts: Dict[str, Union[List[DataLoader[Tuple[Tensor, Tensor]]], str]],
                  loss_fn: Callable[[Tuple[Tensor, Tensor]], Tensor],
                  writer: SummaryWriter,
-                 KFold_class:Optional[BaseCrossValidator]=None) -> None:
+                 training_metric: Optional[Callable[[Tensor, Tensor], float]] = None,
+                 k_fold_class: Optional[BaseCrossValidator] = None) -> None:
         self.models_dicts = models_dicts
         self.dataloaders_dicts = dataloaders_dicts
         self.loss_fn = loss_fn
         self.writer = writer
         if not self.writer:
             warnings.warn("No writer detected")
-        
-        if not KFold_class:
-            self.KFold_class = KFold(5, shuffle=True)
-        else:
-            self.KFold_class = KFold_class
 
+        if not k_fold_class:
+            self.k_fold_class = KFold(5, shuffle=True)
+        else:
+            self.k_fold_class = k_fold_class
+        if training_metric:
+            self.training_metric = training_metric
+        else:
+            self.training_metric = accuracy
         if not isinstance(self.models_dicts, list):
             raise TypeError("The provided models must be a Python list of dictionaries")
 
@@ -67,156 +77,110 @@ def __init__(self,
 
     def start(self,
               optimizer: Type[Optimizer],
-              n_epochs:int=10,
-              cross_validation:bool=False,
-              optimizer_param:Optional[Dict[str, Any]]=None,
-              dataloaders_param:Optional[Dict[str, Any]]=None,
-              lr_scheduler: Optional[Type[_LRScheduler]]=None,
-              scheduler_params:Optional[Dict[str, Any]]=None,
-              profiling:bool=False,
-              parallel_tpu:bool=False,
-              keep_training:bool=False,
-              store_grad_layer_hist:bool=False,
-              n_accumulated_grads:int=0) -> None:
+              n_epochs: int = 10,
+              cross_validation: bool = False,
+              optimizers_param: Optional[Dict[str, Any]] = None,
+              dataloaders_param: Optional[Dict[str, Any]] = None,
+              lr_scheduler: Optional[Type[_LRScheduler]] = None,
+              scheduler_params: Optional[Dict[str, Any]] = None,
+              profiling: bool = False,
+              parallel_tpu: bool = False,
+              keep_training: bool = False,
+              store_grad_layer_hist: bool = False,
+              n_accumulated_grads: int = 0,
+              writer_tag: str = "") -> None:
         """Method to be called when starting the benchmarking
         
         Args:
-            optimizer (torch.optim):
-                a torch optimizers
-            n_epochs (int):
+            optimizer:
+                a torch optimizers class (not the instance)
+            n_epochs:
                 number of training epochs
-            crossvalidation (bool):
+            cross_validation:
                 whether or not to use cross-validation
-            optimizers_param (dict):
+            optimizers_param:
                 dictionary of the optimizers
-                parameters, e.g. `{"lr": 0.001}`
-            dataloaders_param (dict):
+                parameters, e.g. ``{"lr": 0.001}``
+            dataloaders_param:
                 dictionary of the dataloaders
-                parameters, e.g. `{"batch_size": 32}`
-            lr_scheduler (torch.optim):
-                a learning rate scheduler
-            scheduler_params (dict):
+                parameters, e.g. ``{"batch_size": 32}``
+            lr_scheduler:
+                a learning rate scheduler class (not instance)
+            scheduler_params:
                 learning rate scheduler parameters
-            profiling (bool, default=False):
+            profiling:
                 whether or not you want to activate the
                 profiler
-            parallel_tpu (bool):
+            parallel_tpu:
                 boolean value to run the computations
                 on multiple TPUs
-            keep_training (bool):
+            keep_training:
                 This flag allows to restart a training from
                 the existing optimizer as well as the
                 existing model
-            store_grad_layer_hist (bool):
+            store_grad_layer_hist:
                 This flag allows to store the gradients
                 and the layer values in tensorboard for
                 each epoch
-            n_accumulated_grads (int, default=0):
+            n_accumulated_grads:
                 number of accumulated gradients. It is
                 considered only if a positive integer
+            writer_tag:
+                string to be added to the tensorboard items title
         """
 
         print("Benchmarking Started")
-        _benchmarking_param(self._inner_function,
-                            [self.models_dicts,
-                             self.dataloaders_dicts],
-                            optimizer, n_epochs,
-                            cross_validation,
-                            optimizer_param,
-                            dataloaders_param,
-                            lr_scheduler,
-                            scheduler_params,
-                            profiling,
-                            parallel_tpu,
-                            keep_training,
-                            store_grad_layer_hist,
-                            n_accumulated_grads,
-                            writer_tag="")
+        config = TrainerConfig(
+            optimizer,
+            n_epochs,
+            cross_validation,
+            optimizers_param,
+            dataloaders_param,
+            lr_scheduler,
+            scheduler_params,
+            None,
+            profiling,
+            parallel_tpu,
+            keep_training,
+            store_grad_layer_hist,
+            n_accumulated_grads,
+            writer_tag
+        )
 
+        _benchmarking_param(self._inner_function,
+                            (self.models_dicts,
+                             self.dataloaders_dicts),
+                            config)
 
     def _inner_function(self,
-                        model:Dict[str,torch.nn.Module],
-                        dataloaders:Dict[str, List[DataLoader[Tuple[Tensor, Tensor]]]],
-                        optimizer: Type[Optimizer],
-                        n_epochs:int,
-                        cross_validation:bool,
-                        optimizer_param:Dict[str, Any],
-                        dataloaders_param:Dict[str, Any],
-                        lr_scheduler:Optional[Type[_LRScheduler]],
-                        scheduler_params:Dict[str, Any],
-                        profiling:bool,
-                        parallel_tpu:bool,
-                        keep_training:bool,
-                        store_grad_layer_hist:bool,
-                        n_accumulated_grads:int,
-                        writer_tag:str=""):
+                        model: Dict[str, Union[torch.nn.Module, str]],
+                        dataloaders: Dict[str, Union[List[DataLoader[Tuple[Tensor, Tensor]]], str]],
+                        config: TrainerConfig) -> None:
         """private method to run the inner
         function of the benchmark loops
         
         Args:
-            model (dict):
+            model:
                 dictionary defining the model name
                 and actual nn.Module
-            dataloaders (dict):
+            dataloaders:
                 dictionary defining the dataset name
                 and the actual list of dataloaders, e.g.
                 ``[dl_tr, dl_val, dl_ts]``
-            optimizer (torch.optim):
-                a torch optimizer class
-            n_epochs (int):
-                number of training epochs
-            cross_validation (bool):
-                whether or not to use cross-validation
-            optimizers_param (dict):
-                dictionary of the optimizers
-                parameters, e.g. `{"lr": 0.001}`
-            dataloaders_param (dict):
-                dictionary of the dataloaders
-                parameters, e.g. `{"batch_size": 32}`
-            lr_scheduler (torch.optim):
-                a learning rate scheduler class
-            scheduler_params (dict):
-                learning rate scheduler parameters
-            profiling (bool, default=False):
-                whether or not you want to activate the
-                profiler
-            parallel_tpu (bool):
-                boolean value to run the computations
-                on multiple TPUs
-            keep_training (bool):
-                boolean flag to use the same model for
-                further training
-            store_grad_layer_hist (bool):
-                flag to store or not the layer's grads
-                on tensorboard
-            n_accumulated_grads (int):
-                this is the number of accumated grads. It
-                is taken into account only for positive integers
-            writer_tag (str):
-                the tensorboard writer tag
+            config:
+                the configuration class ``TrainerConfig``.
         """
         pipe = Trainer(model["model"], dataloaders["dataloaders"],
-                        self.loss_fn, self.writer, self.KFold_class)
-        writer_tag += "Dataset:" + dataloaders["name"] +"|Model:" + model["name"]
-        pipe.train(optimizer,
-                   n_epochs,
-                   cross_validation,
-                   optimizer_param,
-                   dataloaders_param,
-                   lr_scheduler,
-                   scheduler_params,
-                   None,
-                   profiling,
-                   parallel_tpu,
-                   keep_training,
-                   store_grad_layer_hist,
-                   n_accumulated_grads,
-                   writer_tag)
-
-def _benchmarking_param(fun: Callable[[Any],Any],
-                        arguments:Tuple[Dict[str, torch.nn.Module],
-                                             Dict[str, DataLoader[Tuple[Tensor, Tensor]]]],
-                        *args, **kwargs):  # type: ignore
+                       self.loss_fn, self.writer, self.training_metric,  # type: ignore
+                       self.k_fold_class)
+        config.writer_tag += "Dataset:" + dataloaders["name"] + "|Model:" + model["name"]  # type: ignore
+        pipe.train(**config.to_dict())
+
+
+def _benchmarking_param(fun: Callable[[Any], Any],
+                        arguments: Tuple[Dict[str, Union[torch.nn.Module, str]],
+                                         Dict[str, Union[List[DataLoader[Tuple[Tensor, Tensor]]], str]]],
+                        config: TrainerConfig) -> None:
     """Function to be used as pseudo-decorator for
     benchmarking loops
     
@@ -227,10 +191,8 @@ def _benchmarking_param(fun: Callable[[Any],Any],
             list of arguments to pass to the inner
             function of the wrapper. Expected to receive
             ``models_dicts, dataloaders_dict = arguments``
-        *args (*list):
-            all the args of ``fun``
-        **kwargs (**dict):
-            all the kwargs of ``fun``
+        config:
+            the configuration class ``TrainerConfig``.
 
     """
 
@@ -238,8 +200,7 @@ def _benchmarking_param(fun: Callable[[Any],Any],
     for dataloaders in dataloaders_dicts:
         for model in models_dicts:
             if _are_compatible(model, dataloaders):
-                print("*"*40)
+                print("*" * 40)
                 print(f"Performing Gridsearch on Dataset: {dataloaders['name']}"
                       f", Model: {model['name']}")
-                fun(model, dataloaders, *args, **kwargs)
-
+                fun(model, dataloaders, config)
diff --git a/gdeep/search/hpo.py b/gdeep/search/hpo.py
index e669e04a..4cbc5ef2 100644
--- a/gdeep/search/hpo.py
+++ b/gdeep/search/hpo.py
@@ -2,39 +2,47 @@
 import time
 import warnings
 from itertools import chain, combinations
-from functools import partial
-from typing import Tuple, Any, Dict, Callable, \
-    Type, Optional, List
+from typing import Tuple, Any, Dict, \
+    Type, Optional, List, Union
 
+from typing_extensions import Literal
 import torch
 import optuna
 from optuna.trial import TrialState
 import pandas as pd
 import numpy as np
-from sklearn.model_selection import KFold
-from torch.utils.data.sampler import SubsetRandomSampler
 from torch.utils.tensorboard import SummaryWriter
 from torch.utils.tensorboard.summary import hparams
 import random
 import string
-from torch.optim import *
+from torch.optim import *  # noqa
 import plotly.express as px
-from optuna.pruners import MedianPruner
+from optuna.pruners import MedianPruner, BasePruner
+from optuna.trial._base import BaseTrial
+from optuna.study import BaseStudy as Study
+from torch.utils.data import DataLoader
 from torch.optim import Optimizer
-from torch.optim.lr_scheduler import _LRScheduler
+from torch.optim.lr_scheduler import _LRScheduler  # noqa
 
-from gdeep.utility import _are_compatible, _inner_refactor_scalars, KnownWarningSilencer
+from gdeep.utility import _inner_refactor_scalars, KnownWarningSilencer  # noqa
 from gdeep.trainer import Trainer
 from gdeep.search import Benchmark, _benchmarking_param
 from gdeep.visualisation import plotly2tensor
 from ..utility import save_model_and_optimizer
-from gdeep.utility import DEVICE
+from .hpo_config import HPOConfig
+
+Tensor = torch.Tensor
+Array = np.ndarray
+SEARCH_METRICS = ("loss", "accuracy")
 
 
 class GiottoSummaryWriter(SummaryWriter):
-    def add_hparams(self, hparam_dict, metric_dict,
-                    hparam_domain_discrete=None, run_name=None,
-                    scalars_lists=None, best_not_last=False):
+    def add_hparams(self, hparam_dict: Dict[str, Any],
+                    metric_dict: Dict[str, Any],
+                    hparam_domain_discrete: Optional[Dict[str, List[Any]]] = None,
+                    run_name: str = None,
+                    scalars_lists: Optional[List[List[Tuple[Any, int]]]] = None,
+                    best_not_last: bool = False):
         """Add a set of hyperparameters to be compared in TensorBoard.
 
         Args:
@@ -43,24 +51,27 @@ def add_hparams(self, hparam_dict, metric_dict,
                 name of the hyper parameter and it's corresponding value.
                 The type of the value can be one of `bool`, `string`, `float`,
                 `int`, or `None`.
-            metric_dict (dict):
+            metric_dict :
                 Each key-value pair in the dictionary is the
                 name of the metric and it's corresponding value. Note that the key used
                 here should be unique in the tensorboard record. Otherwise the value
                 you added by ``add_scalar`` will be displayed in hparam plugin. In most
                 cases, this is unwanted.
             hparam_domain_discrete:
-                (Optional[Dict[str, List[Any]]]) A dictionary that
+                A dictionary that
                 contains names of the hyperparameters and all discrete values they can hold
-            run_name (str):
+            run_name :
                 Name of the run, to be included as part of the logdir.
                 If unspecified, will use current timestamp.
-            scalars_lists (list):
+            scalars_lists :
                 The lists for the loss and accuracy plots.
                 This is a list with two lists
                 (one for accuracy and one for the loss).
                 Each one of the inner lists contain the
                 pairs (metric_value, epoch).
+            best_not_last:
+                boolean flag to dcide what value to store in the
+                tensorboard tables for the whole training cycle
 
         Examples::
 
@@ -71,12 +82,12 @@ def add_hparams(self, hparam_dict, metric_dict,
                                   {'hparam/accuracy': 10*i, 'hparam/loss': 10*i})
 
         """
-        torch._C._log_api_usage_once("tensorboard.logging.add_hparams")
+        torch._C._log_api_usage_once("tensorboard.logging.add_hparams")  # type: ignore
         if (not isinstance(hparam_dict, dict)) or (not isinstance(metric_dict, dict)):
             raise TypeError('hparam_dict and metric_dict should be dictionary.')
         exp, ssi, sei = hparams(hparam_dict, metric_dict, hparam_domain_discrete)
         if not run_name:
-            run_name = str(time.time()).replace(":","-")
+            run_name = str(time.time()).replace(":", "-")
         logdir = os.path.join(self._get_file_writer().get_logdir(), run_name)
         with SummaryWriter(log_dir=logdir) as w_hp:
             w_hp.file_writer.add_summary(exp)
@@ -85,12 +96,13 @@ def add_hparams(self, hparam_dict, metric_dict,
             if isinstance(scalars_lists, list) or isinstance(scalars_lists, tuple):
                 scalars_list_loss = scalars_lists[0]
                 if best_not_last:
-                    scalars_list_loss = [x for x in scalars_list_loss[:np.argmin(np.array(scalars_list_loss)[:,0])+1]]
+                    scalars_list_loss = [x for x in
+                                         scalars_list_loss[:np.argmin(np.array(scalars_list_loss)[:, 0]) + 1]]
                 for v, t in scalars_list_loss:
                     w_hp.add_scalar("loss", v, t)
                 scalars_list_acc = scalars_lists[1]
                 if best_not_last:
-                    scalars_list_acc = [x for x in scalars_list_acc[:np.argmax(np.array(scalars_list_acc)[:,0])+1]]
+                    scalars_list_acc = [x for x in scalars_list_acc[:np.argmax(np.array(scalars_list_acc)[:, 0]) + 1]]
                 for v, t in scalars_list_acc:
                     w_hp.add_scalar("accuracy", v, t)
 
@@ -101,66 +113,63 @@ class HyperParameterOptimization(Trainer):
     parameters such as learning rate, optimizer.
 
     Args:
-        obj (either Trainer or Benchmark object):
+        obj :
             either a Trainer or a Bechmark class
             instance
-        search_metric (string):
+        search_metric :
             either ``'loss'`` or ``'accuracy'``
-        n_trials (int):
+        n_trials :
             number of total search trials
-        best_not_last (bool, default False):
+        best_not_last :
             boolean flag that is ``True`` would use
             the best metric over epochs averaged over the folds in CV
             rather than the last value of the metrics
             over the epochs averaged over the folds
-        pruner (optuna.Pruners, default MedianPruner):
+        pruner :
             Instance of an optuna pruner, can be user-defined
-        sampler (optuna.Samplers, default TPESampler):
+        sampler :
             If left unspecified, ``TPESample`` is used during single-objective 
             optimization and ``NSGAIISampler`` during multi-objective optimization
-        db_url (str):
+        db_url:
             name of the database to connect to. For example
             ``mysql+mysqldb://usr:psw@host:port/db_name``
-        study_name (str):
+        study_name:
             name of the optuna study
 
     """
 
     def __init__(self, obj: Trainer,
-                 search_metric:str="loss",
-                 n_trials:int=10,
-                 best_not_last:bool=False,
-                 pruner=None,
+                 search_metric: Literal["loss", "accuracy"] = "loss",
+                 n_trials: int = 10,
+                 best_not_last: bool = False,
+                 pruner: Optional[BasePruner] = None,
                  sampler=None,
-                 db_url:str=None,
-                 study_name:str=None):
+                 db_url: Optional[str] = None,
+                 study_name: Optional[str] = None):
         self.best_not_last_gs = best_not_last
         self.is_pipe = None
-        self.study = None
+        self.study: Optional[Study] = None
         self.best_val_acc_gs = 0
         self.best_val_loss_gs = np.inf
-        self.list_res = []
+        self.list_res: List[Any] = []
         self.df_res = None
         self.db_url = db_url
         self.study_name = study_name
-        if (isinstance(obj, Trainer)):
+        if isinstance(obj, Trainer):
             self.pipe = obj
             super().__init__(self.pipe.model,
                              self.pipe.dataloaders,
                              self.pipe.loss_fn,
                              self.pipe.writer,
-                             self.pipe.KFold_class)
+                             self.pipe.training_metric,
+                             self.pipe.k_fold_class)
             # Pipeline.__init__(self, obj.model, obj.dataloaders, obj.loss_fn, obj.writer)
             self.is_pipe = True
-        elif (isinstance(obj, Benchmark)):
+        elif isinstance(obj, Benchmark):
             self.bench = obj
             self.is_pipe = False
-
-        self.search_metric = (search_metric if search_metric in ("loss", "accuracy")
-                              else None)
-        if self.search_metric is None:
-            raise ValueError("Wrong search_metric! "
-                             "Either `loss` or `accuracy`")
+        self.search_metric = search_metric
+        assert self.search_metric in SEARCH_METRICS, "Wrong search_metric! Either `loss` or `accuracy`"
         self.n_trials = n_trials
         self.val_epoch = 0
         self.train_epoch = 0
@@ -172,11 +181,12 @@ def __init__(self, obj: Trainer,
                                        n_warmup_steps=0,
                                        interval_steps=1,
                                        n_min_trials=1)
-        self.scalars_dict = dict()
+        self.scalars_dict: Dict[str, Any] = dict()
         # can be changed by changing this attribute
-        self.store_pickle = False
+        self.store_pickle: bool = False
 
-    def _initialise_new_model(self, models_hyperparam):
+    def _initialise_new_model(self, models_hyperparam: Dict[str, Any]) \
+            -> torch.nn.Module:
         """private method to find the maximal compatible set
         between models and hyperparameters
         
@@ -188,15 +198,16 @@ def _initialise_new_model(self, models_hyperparam):
             nn.Module
                 torch nn.Module
         """
-        
+
         list_of_params_keys = HyperParameterOptimization._powerset(list(models_hyperparam.keys()))
         list_of_params_keys.reverse()
         for params_keys in list_of_params_keys:
-            sub_models_hyperparam = {k:models_hyperparam[k] for k in models_hyperparam.keys() if k in params_keys}
+            sub_models_hyperparam = {k: models_hyperparam[k]
+                                     for k in models_hyperparam.keys() if k in params_keys}
             try:
-                #print(sub_models_hyperparam)
-                new_model = type(self.model)(**sub_models_hyperparam)
-                #print(new_model.state_dict())
+                # print(sub_models_hyperparam)
+                new_model = type(self.model)(**sub_models_hyperparam)  # noqa
+                # print(new_model.state_dict())
                 raise ValueError
             except TypeError:  # when the parameters do not match the model
                 pass
@@ -204,154 +215,103 @@ def _initialise_new_model(self, models_hyperparam):
                 break
 
         try:
-            new_model
+            new_model  # noqa
         except NameError:
             warnings.warn("Model cannot be re-initialised. Using existing one.")
             new_model = self.model
         return new_model
 
-    def _objective(self, trial,
-                   optimizers,
-                   n_epochs,
-                   cross_validation,
-                   optimizers_params,
-                   dataloaders_params,
-                   models_hyperparams,
-                   lr_scheduler,
-                   schedulers_params,
-                   profiling,
-                   parallel_tpu,
-                   keep_training,
-                   store_grad_layer_hist,
-                   n_accumulated_grads,
-                   writer_tag=""):
+    def _objective(self,
+                   trial: Optional[BaseTrial],
+                   config: HPOConfig):
         """default callback function for optuna's study
         
         Args:
-            trial (optuna.trial):
+            trial:
                 the independent variable
-            optimizers (list of torch.optim):
-                list of torch optimizers
-            n_epochs (int):
-                number of training epochs
-            optimizers_params (dict):
-                dictionary of the optimizers
-                parameters, e.g. `{"lr": 0.001}`
-            dataloaders_hyperparams (dict):
-                dictionary of the dataloaders
-                parameters
-            models_hyperparams (dict):
-                dictionary of the model
-                parameters
-            lr_scheduler (torch.optim):
-                a learning rate scheduler
-            schedulers_params (dict):
-                learning rate scheduler parameters
-            profiling (bool):
-                whether or not you want to activate the
-                profiler
-            parallel_tpu (bool):
-                boolean value to run the computations
-                on multiple TPUs
-            keep_training (bool):
-                This flag allows to restart a training from
-                the existing optimizer as well as the
-                existing model
-            store_grad_layer_hist (bool):
-                This flag allows to store the gradients
-                and the layer values in tensorboard for
-                each epoch
-            n_accumulated_grads (int):
-                this is the number of accumated grads. It
-                is taken into account only for positive integers
-            writer_tag (string):
-                tag to prepend to the ouput
-                on tensorboard
+            config:
+                configuration class HPOConfig,
+                containing all the parameters needed
+
+
+        Returns:
+            float
+                metric (either loss or accuracy)
         """
 
         # for proper storing of data
-        self._cross_validation = cross_validation
-        self._k_folds = self.KFold_class.n_splits
+        self._cross_validation = config.cross_validation
+        self._k_folds = self.k_fold_class.n_splits
         # generate optimizer
-        optimizers_names = list(map(lambda x: x.__name__, optimizers))
+        optimizers_names = list(map(lambda x: x.__name__, config.optimizers))
         optimizer = eval(trial.suggest_categorical("optimizer", optimizers_names))
 
         # generate all the hyperparameters
-        self.optimizers_param = HyperParameterOptimization._suggest_params(trial, optimizers_params)
-        self.dataloaders_param = HyperParameterOptimization._suggest_params(trial, dataloaders_params)
-        self.models_hyperparam = HyperParameterOptimization._suggest_params(trial, models_hyperparams)
-        self.schedulers_param = HyperParameterOptimization._suggest_params(trial, schedulers_params)
+        self.optimizers_param = HyperParameterOptimization._suggest_params(trial, config.optimizers_params)
+        self.dataloaders_param = HyperParameterOptimization._suggest_params(trial, config.dataloaders_params)
+        self.models_hyperparam = HyperParameterOptimization._suggest_params(trial, config.models_hyperparams)
+        self.schedulers_param = HyperParameterOptimization._suggest_params(trial, config.schedulers_params)
         # tag for storing the results
-        writer_tag += "/" + str(trial.datetime_start) # str(self.optimizers_param) + \
-            #str(self.dataloaders_param) + str(self.models_hyperparam) + \
-            #str(self.schedulers_param)
+        config.writer_tag += "/" + str(trial.datetime_start)  # str(self.optimizers_param) + \
+        # str(self.dataloaders_param) + str(self.models_hyperparam) + \
+        # str(self.schedulers_param)
         # create a new model instance
         self.model = self._initialise_new_model(self.models_hyperparam)
         self.pipe = Trainer(self.model, self.dataloaders, self.loss_fn,
-                             self.writer, self.KFold_class)
+                            self.writer, self.training_metric, self.k_fold_class)
         # set best_not_last
         self.pipe.best_not_last = self.best_not_last_gs
         # set the run_name
-        self.pipe.run_name = str(trial.datetime_start).replace(":","-")
-        loss, accuracy = self.pipe.train(optimizer, n_epochs,
-                                         cross_validation,
+        self.pipe.run_name = str(trial.datetime_start).replace(":", "-")
+        loss, accuracy = self.pipe.train(optimizer, config.n_epochs,
+                                         config.cross_validation,
                                          self.optimizers_param,
                                          self.dataloaders_param,
-                                         lr_scheduler,
+                                         config.lr_scheduler,
                                          self.schedulers_param,
                                          (trial, self.search_metric),
-                                         profiling,
-                                         parallel_tpu,
-                                         keep_training,
-                                         store_grad_layer_hist,
-                                         n_accumulated_grads,
-                                         writer_tag
+                                         config.profiling,
+                                         config.parallel_tpu,
+                                         config.keep_training,
+                                         config.store_grad_layer_hist,
+                                         config.n_accumulated_grads,
+                                         config.writer_tag
                                          )
 
         # dict of the loss per epoch of the training here above. If CV, multiple scalars are stored
         scalars_dict_value = [self.pipe.val_loss_list_hparam,
                               self.pipe.val_acc_list_hparam]
-        scalar_dict_key = str(trial.datetime_start).replace(":","-")
+        scalar_dict_key = str(trial.datetime_start).replace(":", "-")
         self.scalars_dict[scalar_dict_key] = scalars_dict_value
         # release the run_name
         self.pipe.run_name = None
-        best_loss = self.pipe.best_val_loss
-        best_accuracy = self.pipe.best_val_acc
         self.writer.flush()
         # print
         self._print_output()
-        
+
         # save model and optimizer
         save_model_and_optimizer(self.pipe.model,
-                                 trial_id=str(trial.datetime_start).replace(":","-"),
+                                 trial_id=str(trial.datetime_start).replace(":", "-"),
                                  optimizer=self.pipe.optimizer,
                                  store_pickle=self.store_pickle)
 
         # save results to tensorboard
-        model_name, dataset_name = self._extract_model_and_dataset_name(writer_tag)
+        model_name, dataset_name = self._extract_model_and_dataset_name(config.writer_tag)
 
         self._store_trial_to_tb(trial, scalars_dict_value, scalar_dict_key, model_name,
                                 dataset_name, loss, accuracy)
 
         # returns
         if self.search_metric == "loss":
-            #if self.best_not_last:
-            #    self.best_val_acc_gs = max(self.best_val_acc_gs, best_accuracy)
-            #    self.best_val_loss_gs = min(self.best_val_loss_gs, best_loss)
-            #    return best_loss
             self.best_val_acc_gs = max(self.best_val_acc_gs, accuracy)
             self.best_val_loss_gs = min(self.best_val_loss_gs, loss)
             return loss
         else:
-            #if self.best_not_last:
-            #    self.best_val_acc_gs = max(self.best_val_acc_gs, best_accuracy)
-            #    self.best_val_loss_gs = min(self.best_val_loss_gs, best_loss)
-            #    return best_accuracy
             self.best_val_acc_gs = max(self.best_val_acc_gs, accuracy)
             self.best_val_loss_gs = min(self.best_val_loss_gs, loss)
             return accuracy
 
-    def _extract_model_and_dataset_name(self, writer_tag):
+    def _extract_model_and_dataset_name(self, writer_tag: str) -> Tuple[str, str]:
         """Extract the model and dataset anme from the writer_tag"""
         index_ds = writer_tag.find("Dataset:")
         if index_ds == -1:
@@ -366,49 +326,55 @@ def _extract_model_and_dataset_name(self, writer_tag):
         return model_name, dataset_name
 
     def start(self,
-              optimizers: List[Optimizer],
-              n_epochs:int=1,
-              cross_validation:bool=False,
-              optimizers_params: Optional[Dict[str, Any]]=None,
-              dataloaders_params: Optional[Dict[str, Any]]=None,
-              models_hyperparams: Optional[Dict[str, Any]]=None,
-              lr_scheduler: Type[_LRScheduler]=None,
-              schedulers_params:Optional[Dict[str, Any]]=None,
-              profiling:bool=False,
-              parallel_tpu:bool=False,
-              keep_training:bool=False,
-              store_grad_layer_hist:bool=False,
-              n_accumulated_grads:int=0,
-              writer_tag=""):
+              optimizers: List[Type[Optimizer]],
+              n_epochs: int = 1,
+              cross_validation: bool = False,
+              optimizers_params: Optional[Dict[str, Any]] = None,
+              dataloaders_params: Optional[Dict[str, Any]] = None,
+              models_hyperparams: Optional[Dict[str, Any]] = None,
+              lr_scheduler: Type[_LRScheduler] = None,
+              schedulers_params: Optional[Dict[str, Any]] = None,
+              profiling: bool = False,
+              parallel_tpu: bool = False,
+              keep_training: bool = False,
+              store_grad_layer_hist: bool = False,
+              n_accumulated_grads: int = 0,
+              writer_tag: str="") -> None:
         """method to be called when starting the gridsearch
 
         Args:
-            optimizers (list of torch.optim):
-                list of torch optimizers
-            n_epochs (int):
+            optimizers:
+                list of torch optimizers classes, not isntances
+            n_epochs:
                 number of training epochs
-            cross_validation (bool):
+            cross_validation:
                 whether or not to use cross-validation
-            optimizers_params (dict):
+            optimizers_params:
                 dictionary of optimizers params
-            dataloaders_params (int):
+            dataloaders_params:
                 dictionary of dataloaders parameters
-            models_hyperparams (dict):
+            models_hyperparams:
                 dictionary of model parameters
-            lr_scheduler (torch.optim):
+            lr_scheduler:
                 torch learning rate schduler class
-            schedulers_params (dict):
+            schedulers_params:
                 learning rate scheduler parameters
-            profiling (bool, default=False):
+            profiling :
                 whether or not you want to activate the
                 profiler
-            parallel_tpu (bool):
+            parallel_tpu:
                 boolean value to run the computations
                 on multiple TPUs
             n_accumulated_grads (int, default=0):
                 number of accumulated gradients. It is
                 considered only if a positive integer
-            writer_tag (str):
+            keep_training:
+                bool flag to decide whether to continue
+                training or not
+            store_grad_layer_hist:
+                flag to store the gradents of the layers in the
+                tensorboard histograms
+            writer_tag:
                 tag to prepend to the ouput
                 on tensorboard
         """
@@ -417,113 +383,78 @@ def start(self,
                                              sampler=self.sampler,
                                              pruner=self.pruner,
                                              storage=self.db_url,
-                                             study_name=self.study_name if self.study_name is not None else \
-                                                 ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(20)),
+                                             study_name=self.study_name if self.study_name is not None else
+                                             ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in
+                                                     range(20)),
                                              load_if_exists=True)
         else:
             self.study = optuna.create_study(direction="maximize",
                                              sampler=self.sampler,
                                              pruner=self.pruner,
                                              storage=self.db_url,
-                                             study_name=self.study_name if self.study_name is not None else \
-                                                 ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in
-                                                         range(20)),
+                                             study_name=self.study_name if self.study_name is not None else
+                                             ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in
+                                                     range(20)),
                                              load_if_exists=True)
+        config = HPOConfig(
+            optimizers,
+            n_epochs,
+            cross_validation,
+            optimizers_params,
+            dataloaders_params,
+            models_hyperparams,
+            lr_scheduler,
+            schedulers_params,
+            profiling,
+            parallel_tpu,
+            keep_training,
+            store_grad_layer_hist,
+            n_accumulated_grads,
+            writer_tag)
+
         if self.is_pipe:
             # in the __init__, self.model and self.dataloaders are
             # already initialised. So they exist also in _objective()
-            self._inner_optimisat_fun(self.model,self.dataloaders,
-                                      optimizers,
-                                      n_epochs,
-                                      cross_validation,
-                                      optimizers_params,
-                                      dataloaders_params,
-                                      models_hyperparams,
-                                      lr_scheduler,
-                                      schedulers_params,
-                                      profiling,
-                                      parallel_tpu,
-                                      keep_training,
-                                      store_grad_layer_hist,
-                                      n_accumulated_grads,
-                                      writer_tag)
+            self._inner_optimisat_fun(self.model,
+                                      self.dataloaders,
+                                      config)
 
         else:
             _benchmarking_param(self._inner_optimisat_fun,
                                 [self.bench.models_dicts,
                                  self.bench.dataloaders_dicts],
-                                optimizers,
-                                n_epochs,
-                                cross_validation,
-                                optimizers_params,
-                                dataloaders_params,
-                                models_hyperparams,
-                                lr_scheduler,
-                                schedulers_params,
-                                profiling,
-                                parallel_tpu,
-                                keep_training,
-                                store_grad_layer_hist,
-                                n_accumulated_grads,
-                                writer_tag="")
+                                config)
         # self._store_to_tensorboard()
 
-
-    def _inner_optimisat_fun(self, model, dataloaders,
-                             optimizers,
-                             n_epochs,
-                             cross_validation,
-                             optimizers_params,
-                             dataloaders_params,
-                             models_hyperparams,
-                             lr_scheduler,
-                             schedulers_params,
-                             profiling,
-                             parallel_tpu,
-                             keep_training,
-                             store_grad_layer_hist,
-                             n_accumulated_grads,
-                             writer_tag=""):
+    def _inner_optimisat_fun(self, model: Union[torch.nn.Module, Dict[str, torch.nn.Module]],
+                             dataloaders: Union[List[DataLoader[Tuple[Tensor, Tensor]]], Dict[str, DataLoader]],
+                             config: HPOConfig) -> None:
         """private method to be decorated with the
         benchmark decorator to have benchmarking
         or simply used as is if no benchmarking is
         needed
         """
-        
-        
+
         try:
-            writer_tag = "Dataset:" + dataloaders["name"] + \
-                "|Model:" + model["name"]
+            config.writer_tag = "Dataset:" + dataloaders["name"] + \
+                         "|Model:" + model["name"]
             super().__init__(model["model"],
                              dataloaders["dataloaders"],
                              self.bench.loss_fn,
                              self.bench.writer,
-                             self.bench.KFold_class)
+                             self.bench.training_metric,
+                             self.bench.k_fold_class)
         except TypeError:
             pass
 
-        self.study.optimize(lambda tr: self._objective(tr,
-                                                       optimizers,
-                                                       n_epochs,
-                                                       cross_validation,
-                                                       optimizers_params,
-                                                       dataloaders_params,
-                                                       models_hyperparams,
-                                                       lr_scheduler,
-                                                       schedulers_params,
-                                                       profiling,
-                                                       parallel_tpu,
-                                                       keep_training,
-                                                       store_grad_layer_hist,
-                                                       n_accumulated_grads,
-                                                       writer_tag),
+        self.study.optimize(lambda tr: self._objective(tr, config),
                             n_trials=self.n_trials,
                             timeout=None)
 
         try:
             self._results(model_name=model["name"],
                           dataset_name=dataloaders["name"])
-            #save_model_and_optimizer(self.pipe.model,
+            # save_model_and_optimizer(self.pipe.model,
             #                         model["name"] +
             #                         str(self.optimizers_param) +
             #                         str(self.dataloaders_param) +
@@ -534,13 +465,12 @@ def _inner_optimisat_fun(self, model, dataloaders,
             try:
                 self._results(model_name=self.pipe.model.__class__.__name__,
                               dataset_name=self.pipe.dataloaders[0].dataset.__class__.__name__)
-                #save_model_and_optimizer(self.pipe.model,
+                # save_model_and_optimizer(self.pipe.model,
                 #                         optimizer=self.pipe.optimizer)
             except AttributeError:
                 self._results()
 
-
-    def _print_output(self):
+    def _print_output(self) -> None:
         """Printing the results of an optimisation"""
         results_string_to_print = (("\nBest Validation loss: " +
                                     str(self.pipe.best_val_loss))
@@ -557,15 +487,18 @@ def _print_output(self):
                            )
         try:
             # print models, metric and hyperparameters
-            print("*" * 20 + " RESULTS " + 20 * "*"+"\n",
+            print("*" * 20 + " RESULTS " + 20 * "*" + "\n",
                   "\nModel: ", self.pipe.model.__class__.__name__,
                   string_to_print)
         except AttributeError:
-            print("*" * 20 + " RESULTS " + 20 * "*"+"\n" +string_to_print)
-
-    def _store_to_list_each_step(self, trial, model_name,
-                                 dataset_name, loss=np.inf,
-                                 accuracy=-1, list_res: list=None) -> list:
+            print("*" * 20 + " RESULTS " + 20 * "*" + "\n" + string_to_print)
+
+    def _store_to_list_each_step(self, trial: BaseTrial,
+                                 model_name: str,
+                                 dataset_name: str,
+                                 loss: float = np.inf,
+                                 accuracy: float = -1.,
+                                 list_res: Optional[List[Any]] = None) -> List[Any]:
         """Private method to store all the HPO parameters
         of one trial
         
@@ -599,20 +532,24 @@ def _store_to_list_each_step(self, trial, model_name,
                 loss = trial.value
             except AttributeError:
                 loss = loss
-            list_res.append([str(trial.datetime_start).replace(":","-"), model_name,
-                                  dataset_name] + temp_list + [loss, -1])
+            list_res.append([str(trial.datetime_start).replace(":", "-"), model_name,
+                             dataset_name] + temp_list + [loss, -1])
         else:
             try:
                 accuracy = trial.value
             except AttributeError:
                 accuracy = accuracy
-            list_res.append([str(trial.datetime_start).replace(":","-"), model_name,
-                                  dataset_name] + temp_list + [np.inf, accuracy])
+            list_res.append([str(trial.datetime_start).replace(":", "-"), model_name,
+                             dataset_name] + temp_list + [np.inf, accuracy])
         return list_res
 
-    def _store_trial_to_tb(self, trial, scalars_dict_value,
-                           scalar_dict_key, model_name,
-                           dataset_name, loss, accuracy):
+    def _store_trial_to_tb(self, trial: BaseTrial,
+                           scalars_dict_value: Tuple[List[Any], List[Any]],
+                           scalar_dict_key: str,
+                           model_name: str,
+                           dataset_name: str,
+                           loss: float,
+                           accuracy: float) -> None:
         """store hyperparameters of a single trial to
         tensorboard
         """
@@ -620,12 +557,13 @@ def _store_trial_to_tb(self, trial, scalars_dict_value,
                                                  dataset_name, loss, accuracy)
 
         keys = ["model", "dataset"] + \
-               list(trial.params.keys())
+            list(trial.params.keys())
 
         # average over the scalars_dict
         scalars_dict_avg = self._refactor_scalars(scalars_dict_value)
         # dict of parameters
-        dictio = {k: (int(v) if isinstance(v, np.int64) or isinstance(v, np.int32) else v) for k, v in zip(keys, list_res[0][1:-2])}
+        dictio = {k: (int(v) if isinstance(v, np.int64) or isinstance(v, np.int32) else v) for k, v in
+                  zip(keys, list_res[0][1:-2])}
 
         try:
             self.writer.add_hparams(dictio,
@@ -639,20 +577,17 @@ def _store_trial_to_tb(self, trial, scalars_dict_value,
 
         self.writer.flush()
 
-
-
-    def _results(self, model_name="model", dataset_name="dataset"):
+    def _results(self, model_name: str = "model",
+                 dataset_name: str = "dataset") -> pd.DataFrame:
         """This method returns the dataframe with all the results of
         the hyperparameters optimisaton.
         It also saves the figures in the writer.
 
         Args:
-            run_name (str):
-                name of the tensorboard run
-            model_name (str):
+            model_name:
                 the model name for the
                 tensorboard gridsearch table
-            dataset_name (str)
+            dataset_name:
                 the dataset name for the
                 tensorboard gridsearch table
 
@@ -675,25 +610,26 @@ def _results(self, model_name="model", dataset_name="dataset"):
             print("Parameters Values for best trial: ", trial_best.params)
             print("DateTime start of the best trial: ", trial_best.datetime_start)
         except ValueError:
-            warnings.warn("No best trial found.")
-        
+            warnings.warn("No best trial found. Using the first trial")
+            trial_best = self.study.trials[0]
+
         for tria in trials:
             self._store_to_list_each_step(tria, model_name, dataset_name,
                                           tria.value, tria.value, self.list_res)
 
         self.df_res = pd.DataFrame(self.list_res, columns=["run_name", "model", "dataset"] +
-                              list(trial_best.params.keys())+["loss", "accuracy"])
+                                   list(trial_best.params.keys()) + ["loss", "accuracy"])
         # compute hyperparams correlaton
         corr, labels = self._correlation_of_hyperparams()
         if self.n_trials > 1:
             try:
                 fig2 = px.imshow(corr,
-                        labels=dict(x="Parameters",
-                                    y="Parameters",
-                                    color="Correlation"),
-                        x=labels,
-                        y=labels
-                       )
+                                 labels=dict(x="Parameters",
+                                             y="Parameters",
+                                             color="Correlation"),
+                                 x=labels,
+                                 y=labels
+                                 )
                 fig2.update_xaxes(side="top")
                 fig2.show()
                 img2 = plotly2tensor(fig2)
@@ -705,28 +641,28 @@ def _results(self, model_name="model", dataset_name="dataset"):
             except ValueError:
                 warnings.warn("Cannot send the picture of the correlation" +
                               " matrix to tensorboard")
-        
+
         return self.df_res
 
-    def _correlation_of_hyperparams(self):
+    def _correlation_of_hyperparams(self) -> Tuple[Array, List[Any]]:
         """Correlations of numerical hyperparameters"""
         list_of_arrays = []
         labels = []
         for col in self.df_res.columns:
             vals = self.df_res[col].values
-            #print("type here", vals.dtype)
-            if vals.dtype in [np.float16, np.float64, 
+            # print("type here", vals.dtype)
+            if vals.dtype in [np.float16, np.float64,
                               np.float32,
                               np.int32, np.int64,
                               np.int16]:
                 list_of_arrays.append(vals)
                 labels.append(col)
-        #print(list_of_arrays)
-        with KnownWarningSilencer() as var1:
+        # print(list_of_arrays)
+        with KnownWarningSilencer():
             corr = np.corrcoef(np.array(list_of_arrays))
         return corr, labels
 
-    #def _store_to_tensorboard(self):
+    # def _store_to_tensorboard(self):
     #    """Store the hyperparameters to tensorboard"""
     #    # average over the scalars_dict
     #    scalars_dict_avg = dict()
@@ -748,13 +684,14 @@ def _correlation_of_hyperparams(self):
     #
     #    return self.df_res
 
-    def _refactor_scalars(self, two_lists):
+    def _refactor_scalars(self, two_lists: Tuple[List[Any], List[Any]]) \
+            -> Tuple[List[Any], List[Any]]:
         """private method to transform a list with
         many values for the same epoch into a dictionary
         compatible with ``add_scalar`` averaged per epoch
 
         Args:
-            two_lists (list):
+            two_lists :
                 two lists with pairs (value, time) with possible
                 repetition of the same time
 
@@ -764,18 +701,16 @@ def _refactor_scalars(self, two_lists):
         """
 
         out0 = _inner_refactor_scalars(two_lists[0],
-                                            self._cross_validation,
-                                            self._k_folds)
+                                       self._cross_validation,
+                                       self._k_folds)
         out1 = _inner_refactor_scalars(two_lists[1],
-                                            self._cross_validation,
-                                            self._k_folds)
+                                       self._cross_validation,
+                                       self._k_folds)
         return out0, out1
 
-
-
-
     @staticmethod
-    def _suggest_params(trial, params):
+    def _suggest_params(trial: BaseTrial,
+                        params: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]:
         """Utility function to generate the parameters
         for the gridsearch. It is based on optuna `suggest_<type>`.
 
@@ -792,21 +727,21 @@ def _suggest_params(trial, params):
         if params is None:
             return None
         for k, v in params.items():
-            if (isinstance(v, list) or isinstance(v, tuple)) and len(v)==1:
-                params[k] = 2*v
+            if (isinstance(v, list) or isinstance(v, tuple)) and len(v) == 1:
+                params[k] = 2 * v
         param_temp = {}
-        param_temp = {k:HyperParameterOptimization._new_suggest_float(trial, k,*v) for k,v in
+        param_temp = {k: HyperParameterOptimization._new_suggest_float(trial, k, *v) for k, v in
                       params.items() if (isinstance(v, list) or isinstance(v, tuple))
                       and (isinstance(v[0], float) or isinstance(v[1], float))}
-        param_temp2 = {k:trial.suggest_int(k,*v) for k,v in
+        param_temp2 = {k: trial.suggest_int(k, *v) for k, v in
                        params.items() if (isinstance(v, list) or isinstance(v, tuple))
                        and (isinstance(v[0], int) or isinstance(v[1], int))}
         param_temp.update(param_temp2)
-        param = {k:trial.suggest_categorical(k, v) for k,v in
+        param = {k: trial.suggest_categorical(k, v) for k, v in
                  params.items() if (isinstance(v, list) or isinstance(v, tuple))
                  and (isinstance(v[0], str) or isinstance(v[1], str))}
         param.update(param_temp)
-        #print(param)
+        # print(param)
         return param
 
     @staticmethod
@@ -816,7 +751,6 @@ def _new_suggest_float(trial, name, low, high, step=None, log=False):
 
     @staticmethod
     def _powerset(iterable):
-        "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
+        """powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"""
         s = list(iterable)
         return list(chain.from_iterable(combinations(s, r) for r in range(len(s) + 1)))
-
diff --git a/gdeep/search/hpo_config.py b/gdeep/search/hpo_config.py
new file mode 100644
index 00000000..a76eda68
--- /dev/null
+++ b/gdeep/search/hpo_config.py
@@ -0,0 +1,66 @@
+from dataclasses import dataclass
+from typing import List, Any, Type, Optional, Dict
+
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler  # noqa
+
+
+@dataclass
+class HPOConfig:
+    """Config class for thee HPO parameters
+
+    Args:
+        optimizers:
+            list of torch optimizers classes, not isntances
+        n_epochs:
+            number of training epochs
+        cross_validation:
+            whether or not to use cross-validation
+        optimizers_params:
+            dictionary of optimizers params
+        dataloaders_params:
+            dictionary of dataloaders parameters
+        models_hyperparams:
+            dictionary of model parameters
+        lr_scheduler:
+            torch learning rate schduler class
+        schedulers_params:
+            learning rate scheduler parameters
+        profiling :
+            whether or not you want to activate the
+            profiler
+        parallel_tpu:
+            boolean value to run the computations
+            on multiple TPUs
+        n_accumulated_grads (int, default=0):
+            number of accumulated gradients. It is
+            considered only if a positive integer
+        keep_training:
+            bool flag to decide whether to continue
+            training or not
+        store_grad_layer_hist:
+            flag to store the gradents of the layers in the
+            tensorboard histograms
+        writer_tag:
+            tag to prepend to the ouput
+            on tensorboard"""
+    optimizers: List[Type[Optimizer]]
+    n_epochs: int = 1
+    cross_validation: bool = False
+    optimizers_params: Optional[Dict[str, Any]] = None
+    dataloaders_params: Optional[Dict[str, Any]] = None
+    models_hyperparams: Optional[Dict[str, Any]] = None
+    lr_scheduler: Type[_LRScheduler] = None
+    schedulers_params: Optional[Dict[str, Any]] = None
+    profiling: bool = False
+    parallel_tpu: bool = False
+    keep_training: bool = False
+    store_grad_layer_hist: bool = False
+    n_accumulated_grads: int = 0
+    writer_tag: str = ""
+
+    def to_dict(self) -> Dict[str, Any]:
+        """method to transform the config file into a dictionary"""
+        return {
+            k: v for k, v in self.__dict__.items() if not k.startswith("_")
+        }
diff --git a/gdeep/search/persformer_hyperparamter_search.py b/gdeep/search/persformer_hyperparamter_search.py
index a686c8df..701fd327 100644
--- a/gdeep/search/persformer_hyperparamter_search.py
+++ b/gdeep/search/persformer_hyperparamter_search.py
@@ -1,13 +1,10 @@
 from sympy import N
 import torch
-from torch.optim import AdamW, Adam, SGD
-from transformers import get_cosine_with_hard_restarts_schedule_with_warmup # type: ignore
+from torch.optim import *
+from transformers import get_cosine_with_hard_restarts_schedule_with_warmup  # type: ignore
 from sklearn.model_selection import StratifiedKFold, KFold
-
 from dotmap import DotMap  # type:ignore
-
 from json import load
-
 from optuna.pruners import NopPruner
 
 from gdeep.topology_layers import Persformer
@@ -15,6 +12,7 @@
 from gdeep.data.datasets import DlBuilderFromDataCloud
 from gdeep.trainer import Trainer
 
+
 class PersformerHyperparameterSearch:
     """This class is used to perform hyperparameter search for Persfomer using 
     cross-validation. The search is performed using the Giotto-Deep GridSearch 
@@ -24,34 +22,35 @@ class PersformerHyperparameterSearch:
     saved in the path provided in the constructor.
     
     Args:
-        dataset_name (str):
+        dataset_name :
             name of the dataset to be used for the 
             hyperparameter search. The dataset must be present in the
             DatasetCloud.
-        download_directory (str):
+        download_directory :
             directory where the dataset is
             either downloaded or already present
-        path_hpo_metadata (str):
+        path_hpo_metadata :
             path to the metadata file containing the
             hyperparameter dictionaries specifying the search space and the
             search metric
-        path_writer (str):
+        path_writer :
             path to the Tensorflow writer directory where the
             Tensorflow summaries are saved during the search process.
             
     Returns:
         None
-    """ 
+    """
+
     def __init__(self,
-                dataset_name,
-                download_directory,
-                path_hpo_metadata,
-                path_writer) -> None:
+                 dataset_name: str,
+                 download_directory: str,
+                 path_hpo_metadata: str,
+                 path_writer: str) -> None:
         self.dataset_name = dataset_name
         self.download_directory = download_directory
         self.path_hpo_metadata = path_hpo_metadata
         self.path_writer = path_writer
-        
+
     def _get_data_loader(self) -> torch.utils.data.DataLoader:
         """Returns the data loader for the dataset specified in the constructor.
         
@@ -60,14 +59,13 @@ def _get_data_loader(self) -> torch.utils.data.DataLoader:
             data loader for the dataset specified in the constructor
         """
         dl_cloud_builder = DlBuilderFromDataCloud(self.dataset_name,
-                                   self.download_directory)
+                                                  self.download_directory)
 
         # create the dataset from the downloaded dataset
-        train_dataloader, _, _ = dl_cloud_builder\
-            .build(({"batch_size" : 10},))
+        train_dataloader, _, _ = dl_cloud_builder \
+            .build([{"batch_size": 10}])
         return train_dataloader
-    
-    
+
     def search(self) -> None:
         """Performs the hyperparameter search. The search is performed using
         the Giotto-Deep GridSearch class. The hyperparameter dictionaries
@@ -79,18 +77,17 @@ def search(self) -> None:
         Returns:
             None
         """
-        
+
         model = Persformer()
-        
+
         train_dataloader = self._get_data_loader()
-        
 
         # initialize loss
         loss_fn = torch.nn.CrossEntropyLoss()
 
         # Initialize the Tensorflow writer
         writer = GiottoSummaryWriter(self.path_writer)
-        
+
         # load hpo metadata
         with open(self.path_hpo_metadata) as config_data_file:
             hyperparameters_dicts = DotMap(load(config_data_file))
@@ -100,15 +97,14 @@ def search(self) -> None:
             schedulers_params = hyperparameters_dicts.schedulers_params
 
         # Initialize pipeline        
-        pipe = Trainer(model, [train_dataloader, None], loss_fn, writer,
-                            eval(
-                                hyperparameters_dicts.fold_mode + "(" +
-                                str(hyperparameters_dicts.n_splits) + ", "+
-                                "shuffle=" + str(hyperparameters_dicts.shuffle)
-                                + ")"
-                            )
-                        )
-
+        pipe = Trainer(model, [train_dataloader, None], loss_fn, writer, None,
+                       eval(
+                           hyperparameters_dicts.fold_mode + "(" +
+                           str(hyperparameters_dicts.n_splits) + ", " +
+                           "shuffle=" + str(hyperparameters_dicts.shuffle)
+                           + ")"
+                       )
+                       )
 
         pruner = NopPruner()
         search = HyperParameterOptimization(pipe,
@@ -119,13 +115,13 @@ def search(self) -> None:
 
         # starting the hyperparameter search
         search.start([eval(opt) for opt in hyperparameters_dicts.optimizer],
-                    n_epochs=schedulers_params.num_training_steps[0],
-                    cross_validation=hyperparameters_dicts.cross_validation,
-                    optimizers_params=optimizers_params,
-                    dataloaders_params=dataloaders_params,
-                    models_hyperparams=models_hyperparams,
-                    lr_scheduler=eval(hyperparameters_dicts.scheduler[0]),
-                    schedulers_params=schedulers_params)
+                     n_epochs=schedulers_params.num_training_steps[0],
+                     cross_validation=hyperparameters_dicts.cross_validation,
+                     optimizers_params=optimizers_params,
+                     dataloaders_params=dataloaders_params,
+                     models_hyperparams=models_hyperparams,
+                     lr_scheduler=eval(hyperparameters_dicts.scheduler[0]),
+                     schedulers_params=schedulers_params)
 
         # Close the Tensorflow writer
-        writer.close()
\ No newline at end of file
+        writer.close()
diff --git a/gdeep/trainer/__init__.py b/gdeep/trainer/__init__.py
index 1b75a015..9ebf9502 100644
--- a/gdeep/trainer/__init__.py
+++ b/gdeep/trainer/__init__.py
@@ -1,8 +1,10 @@
-
+from .metrics import accuracy
 from .trainer import Trainer
-# from .gridsearch import GridSearch
+from .trainer_config import TrainerConfig
 
 
 __all__ = [
-    'Trainer'
+    'Trainer',
+    'accuracy',
+    'TrainerConfig'
     ]
diff --git a/gdeep/trainer/metrics.py b/gdeep/trainer/metrics.py
new file mode 100644
index 00000000..f370058a
--- /dev/null
+++ b/gdeep/trainer/metrics.py
@@ -0,0 +1,29 @@
+
+import torch
+
+Tensor = torch.Tensor
+
+
+def accuracy(prediction: Tensor, y: Tensor) -> float:
+    """This function computes the accuracy
+    for the given prediction and expected
+    output
+
+    Args:
+        prediction:
+            the output of your model. If ``X`` is
+            a tensor, then ``prediction = model(X)``
+        y:
+            the corresponding expected results
+
+    Returns:
+        float:
+            the value of the accuracy
+        """
+    correct: float = 0.
+    try:
+        correct += (prediction.argmax(1) == y).to(torch.float).sum().item()
+    except RuntimeError:
+        correct += (prediction.argmax(2) == y).to(torch.float).sum().item()
+
+    return correct/y.shape[0]*100
diff --git a/gdeep/trainer/trainer.py b/gdeep/trainer/trainer.py
index 22148209..bb66aaf3 100644
--- a/gdeep/trainer/trainer.py
+++ b/gdeep/trainer/trainer.py
@@ -3,10 +3,10 @@
 import time
 from functools import wraps
 import warnings
-from typing import Tuple, Union, Optional, Callable, \
+from typing import Tuple, Optional, Callable, \
     Any, Dict, List, Type
 
-import torch.nn.functional as F
+import torch.nn.functional as f
 from torch.optim import Optimizer
 import torch
 from optuna.trial._base import BaseTrial
@@ -19,12 +19,13 @@
 import optuna
 from datetime import datetime
 from torch.utils.tensorboard import SummaryWriter
-from torch.optim.lr_scheduler import _LRScheduler
+from torch.optim.lr_scheduler import _LRScheduler  # noqa
 
 from ..utility.optimisation import MissingClosureError
 from gdeep.models import ModelExtractor
 from gdeep.utility import _inner_refactor_scalars
 from gdeep.utility import DEVICE
+from .metrics import accuracy
 
 Tensor = torch.Tensor
 
@@ -33,35 +34,34 @@
     import torch_xla.core.xla_model as xm
     import torch_xla.distributed.xla_multiprocessing as xmp
     import torch_xla.distributed.parallel_loader as pl
+
     print("Using TPU!")
-except:
+except ModuleNotFoundError:
     print("No TPUs...")
 
 
-def _add_data_to_tb(func: Callable[[Any], Any]) -> Callable[[Any],Any]:
+def _add_data_to_tb(func: Callable[[Any], Any]) -> Callable[[Any], Any]:
     """decorator to store PR data to tensorboard"""
+
     @wraps(func)
     def wrapper(*args, **kwargs):
         pred, val_loss, correct = func(*args, **kwargs)
         try:
             # add data to tensorboard
             Trainer._add_pr_curve_tb(torch.vstack(pred), kwargs["class_label"],
-                                      kwargs["class_probs"], kwargs["writer"],
-                                      kwargs["writer_tag"] + "/validation")
-            try:
-                kwargs["writer"].flush()
-            except AttributeError:
-                pass
+                                     kwargs["class_probs"],
+                                     kwargs["writer_tag"] + "/validation")
         except NotImplementedError:
             warnings.warn("The PR curve is not being filled because too few data exist")
         return pred, val_loss, correct
+
     return wrapper
-    
-    
+
+
 class Trainer:
     """This is the generic class that allows
-    the user to benchhmark models over architectures
-    datasets, regularisations, metrics... in one line
+    the user to benchmark models over architectures
+    datasets, regularisation, metrics... in one line
     of code.
 
     Args:
@@ -74,7 +74,11 @@ class Trainer:
             loss function to average over batches
         writer :
             tensorboard writer
-        KFold_class (sklearn.model_selection, default ``KFold(5, shuffle=True)``):
+        training_metric:
+            the function that computes the metric: it shall
+            have two arguments, one for the prediction
+            and the other for the ground truth
+        k_fold_class (sklearn.model_selection, default ``KFold(5, shuffle=True)``):
             the class instance to implement the KFold, can be
             any of the Splitter classes of sklearn. More
             info at https://scikit-learn.org/stable/modules/classes.html#module-sklearn.model_selection
@@ -109,52 +113,54 @@ def forward(self, x):
         writer = GiottoSummaryWriter()
         # pipeline
         pipe = Trainer(model, [dl_tr, dl_val, dl_ts],
-                        loss_fn, writer,
+                        loss_fn, writer, None,
                         StratifiedKFold(5, shuffle=True))
         # then one needs to train the model using the pipeline!
         pipe.train(SGD, 2, True, {"lr": 0.001}, n_accumulated_grads=5)
 
     """
 
-    # def __init__(self, model, dataloaders, loss_fn, writer,
-    # hyperparams_search = False, search_metric = "accuracy", n_trials = 10):
     def __init__(self, model: torch.nn.Module,
                  dataloaders: List[DataLoader[Tuple[Tensor, Tensor]]],
-                 loss_fn: Callable[[Tensor,Tensor],Tensor],
-                 writer: SummaryWriter,
-                 KFold_class:Optional[BaseCrossValidator]=None) -> None:
+                 loss_fn: Callable[[Tensor, Tensor], Tensor],
+                 writer: Optional[SummaryWriter] = None,
+                 training_metric: Optional[Callable[[Tensor, Tensor], float]] = None,
+                 k_fold_class: Optional[BaseCrossValidator] = None) -> None:
         self.model = model
         self.initial_model = copy.deepcopy(self.model)
-        assert len(dataloaders) > 0 and len(dataloaders) < 4, "Length of dataloaders must be 1, 2, or 3"
+        assert 0 < len(dataloaders) < 4, "Length of dataloaders must be 1, 2, or 3"
         self.dataloaders = dataloaders  # train and test
         self.train_epoch = 0
         self.val_epoch = 0
         self.best_val_loss = np.inf
         self.best_val_acc = 0
         self.loss_fn = loss_fn
+        if training_metric:
+            self.training_metric = training_metric
+        else:
+            self.training_metric = accuracy
         # integrate tensorboard
         self.writer = writer
         if self.writer is None:
             warnings.warn("No writer detected")
         self.DEVICE = DEVICE
-        self.check_has_trained:bool = False
+        self.check_has_trained: bool = False
         self.registered_hook = None
         # used in gradient clipping
-        self.clip:int = 5
-        # used by gridsearch:
-        self.run_name:Optional[str] = None
-        self.val_loss_list_hparam:List[float] = []
-        self.val_acc_list_hparam:List[float] = []
-        self.best_not_last:bool = False
+        self.clip: int = 5
+        # used by hpo:
+        self.run_name: Optional[str] = None
+        self.val_loss_list_hparam: List[float] = []
+        self.val_acc_list_hparam: List[float] = []
+        self.best_not_last: bool = False
         # profiler
-        self.prof:Any = None
+        self.prof: Any = None
 
-        if not KFold_class:
-            self.KFold_class = KFold(5, shuffle=True)
+        if not k_fold_class:
+            self.k_fold_class = KFold(5, shuffle=True)
         else:
-            self.KFold_class = KFold_class
+            self.k_fold_class = k_fold_class
 
-        
     def _set_initial_model(self) -> None:
         """This private method is used to set
         the initial_model"""
@@ -168,11 +174,10 @@ def _reset_model(self) -> None:
         self.model = copy.deepcopy(self.initial_model)
 
     def _optimisation_step(self,
-                           dl_tr:DataLoader[Tuple[Tensor, Tensor]],
                            steps: int,
                            loss: Tensor,
-                           t_loss: float,
-                           correct: float,
+                           epoch_loss: float,
+                           batch_metric: float,
                            batch: int,
                            closure: Callable[[], Tensor]) -> float:
         """Backpropagation"""
@@ -185,9 +190,9 @@ def _optimisation_step(self,
                 try:
                     torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
                     self.optimizer.step()
-                except (MissingClosureError, ):
-                    self.optimizer.step(closure)
-                    
+                except (MissingClosureError,):
+                    self.optimizer.step(closure)  # type: ignore
+
         else:  # if we use gradient accumulation techniques
             (loss / self.n_accumulated_grads).backward()
             if (batch + 1) % self.n_accumulated_grads == 0:  # do the optimization step only after the accumulations
@@ -197,31 +202,31 @@ def _optimisation_step(self,
                     try:
                         torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
                         self.optimizer.step()
-                    except (MissingClosureError, ):
-                        self.optimizer.step(closure)
+                    except (MissingClosureError,):
+                        self.optimizer.step(closure)  # type: ignore
                 self.optimizer.zero_grad()
         if batch % 1 == 0:
-            t_loss += loss.item()
-            print("Batch training loss: ", t_loss/(batch+1), " \tBatch training accuracy: ",
-                  correct/((batch+1)*dl_tr.batch_size)*100,
-                  " \t[",batch+1,"/", steps,"]                     ", end='\r')
-        return t_loss
-        
-    def _inner_train_loop(self, 
+            epoch_loss += loss.item()
+            print(f"Batch training loss:  {epoch_loss / (batch + 1)}",
+                  f" \tBatch training {self.training_metric.__name__}: ",
+                  batch_metric,
+                  " \t[", batch + 1, "/", steps, "]                     ", end='\r')
+        return epoch_loss
+
+    def _inner_train_loop(self,
                           dl_tr: DataLoader[Tuple[Tensor, Tensor]],
                           writer_tag: str,
                           size: int,
-                          steps: int,
-                          loss: Tensor,
-                          t_loss: float,
-                          correct: float) -> Tuple[float, float]:
+                          steps: int) -> Tuple[float, float]:
         """Private method to run the loop
         over the batches for the optimisation"""
 
-        if not self.prof is None:
+        if self.prof is not None:
             self.prof.start()
+        metric_list = []
+        epoch_loss = 0.
         for batch, (X, y) in enumerate(dl_tr):
-            def closure():
+            def closure() -> Tensor:
                 loss2 = self.loss_fn(self.model(X), y)
                 loss2.backward()
                 return loss2
@@ -230,23 +235,21 @@ def closure():
             y = y.to(self.DEVICE)
             # Compute prediction and loss
             pred = self.model(X)
-            try:
-                correct += (pred.argmax(1) == y).to(torch.float).sum().item()
-            except RuntimeError:
-                correct += (pred.argmax(2) == y).to(torch.float).sum().item()
+            batch_metric = self.training_metric(pred, y)
+            metric_list.append(batch_metric)
             loss = self.loss_fn(pred, y)
             # Save to tensorboard
             try:
                 self.writer.add_scalar(writer_tag + "/loss/train",
                                        loss.item(),
-                                       self.train_epoch*len(dl_tr) + batch)
+                                       self.train_epoch * len(dl_tr) + batch)
 
                 try:
                     top2_pred = torch.topk(pred, 2, -1).values
                     try:
                         self.writer.add_histogram(writer_tag + "/predictions/train",
                                                   torch.abs(torch.diff(top2_pred, dim=-1)),
-                                                  self.train_epoch*steps + batch)
+                                                  self.train_epoch * steps + batch)
                     except ValueError:
                         warnings.warn(f"The histogram is empty, most likely because your loss"
                                       f" is exploding. Try use gradient clipping.")
@@ -254,24 +257,23 @@ def closure():
                     pass
             except AttributeError:
                 pass
-            t_loss = self._optimisation_step(dl_tr, steps, loss, 
-                                    t_loss, correct, batch, closure)
-        
-            if not self.prof is None:
+            epoch_loss = self._optimisation_step(steps, loss, epoch_loss,
+                                                 batch_metric, batch, closure)
+
+            if self.prof is not None:
                 self.prof.step()
 
-        if not self.prof is None:
+        if self.prof is not None:
             self.prof.stop()
 
-        # accuracy:
-        correct /= size
-        t_loss /= steps
-        return correct, t_loss
-
+        # epoch metric and loss:
+        epoch_metric = sum(metric_list)/len(metric_list)
+        epoch_loss /= steps
+        return epoch_metric, epoch_loss
 
     def _train_loop(self,
                     dl_tr: DataLoader[Tuple[Tensor, Tensor]],
-                    writer_tag:str="") -> Tuple[float, float]:
+                    writer_tag: str = "") -> Tuple[float, float]:
         """private method to run a single training
         loop
         """
@@ -282,38 +284,33 @@ def _train_loop(self,
         self.model = self.model.to(self.DEVICE)
         self.model.train()
         try:
-            size = len(dl_tr.sampler.indices)
+            length: int = len(dl_tr.sampler.indices)  # type: ignore
         except AttributeError:
-            size = len(dl_tr.dataset)
+            length: int = len(dl_tr.dataset)  # type: ignore
         steps = len(dl_tr)
-        loss = torch.tensor(0)
-        correct = 0
-        t_loss = 0
         tik = time.time()
         assert self.n_accumulated_grads <= steps, "The number of" + \
                                                   " accumulated gradients shall be diminished!"
-        correct, t_loss = self._inner_train_loop(dl_tr, 
-                                                 writer_tag,
-                                                 size,
-                                                 steps,
-                                                 loss,
-                                                 t_loss,
-                                                 correct)
-        print(f"Epoch training loss: {t_loss:>8f} \tEpoch training accuracy: {(correct*100):.2f}% ".ljust(100))
+        epoch_metric, epoch_loss = self._inner_train_loop(dl_tr,
+                                                          writer_tag,
+                                                          length,
+                                                          steps)
+        print(f"Epoch training loss: {epoch_loss:>8f} \tEpoch training "
+              f"{self.training_metric.__name__}: {epoch_metric:.2f}% ".ljust(100))
         try:
             self.writer.flush()
         except AttributeError:
             pass
-        print(f"Time taken for this epoch: {round(time.time()-tik):.2f}s")
+        print(f"Time taken for this epoch: {round(time.time() - tik):.2f}s")
         try:
             print(f"Learning rate value: {self.optimizer.param_groups[0]['lr']:0.8f}")
         except KeyError:
             pass
-        return t_loss, correct*100
-    
+        return epoch_metric, epoch_loss
+
     def _val_loop(self,
                   dl_val: DataLoader[Tuple[Tensor, Tensor]],
-                  writer_tag:str="") -> Tuple[float, float]:
+                  writer_tag: str = "") -> Tuple[float, float]:
         """private method to run a single validation
         loop
         """
@@ -323,36 +320,30 @@ def _val_loop(self,
             pass
         self.model = self.model.to(self.DEVICE)
         try:
-            size = len(dl_val.sampler.indices)
+            size = len(dl_val.sampler.indices)  # type: ignore
         except AttributeError:
-            size = len(dl_val.dataset)
-        val_loss, correct = 0., 0.
-        class_label:List[Tensor] = []
-        class_probs:List[List[Tensor]] = []
+            size = len(dl_val.dataset)  # type: ignore
+        class_label: List[Tensor] = []
+        class_probs: List[List[Tensor]] = []
         self.model.eval()
 
-        pred, val_loss, correct = self._inner_loop(dl=dl_val,
-                                                   class_probs=class_probs,
-                                                   class_label=class_label,
-                                                   loss=val_loss,
-                                                   correct=correct,
-                                                   writer=self.writer,
-                                                   writer_tag=writer_tag)
+        pred_list, epoch_loss, epoch_metric = self._inner_loop(dl=dl_val,
+                                                               class_probs=class_probs,
+                                                               class_label=class_label,
+                                                               writer_tag=writer_tag)
         # accuracy
-        correct /= size
-        val_loss /= len(dl_val)
         try:
             if not self.run_name:
-                self.writer.add_scalar(writer_tag + "/accuracy/validation",
-                                       correct, self.val_epoch)
+                self.writer.add_scalar(writer_tag + "/metric/validation",
+                                       epoch_metric, self.val_epoch)
             else:
-                self.val_acc_list_hparam.append([correct,
-                                                 self.val_epoch])
-                self.val_loss_list_hparam.append([val_loss,
+                self.val_acc_list_hparam.append([epoch_metric,
                                                  self.val_epoch])
+                self.val_loss_list_hparam.append([epoch_loss,
+                                                  self.val_epoch])
 
             try:
-                top2_pred = torch.topk(torch.vstack(pred), 2, -1).values
+                top2_pred = torch.topk(torch.vstack(pred_list), 2, -1).values
                 try:
                     self.writer.add_histogram(writer_tag + "/predictions/validation",
                                               torch.abs(torch.diff(top2_pred, dim=-1)),
@@ -364,109 +355,108 @@ def _val_loop(self,
                 pass
         except AttributeError:
             pass
-        print(f"Validation results: \n Accuracy: {(100*correct):.2f}%, \
-                Avg loss: {val_loss:>8f} \n")
+        print(f"Validation results: \n {self.training_metric.__name__}: {epoch_metric:.2f}%, \
+                Avg loss: {epoch_loss:>8f} \n")
         try:
             self.writer.flush()
         except AttributeError:
             pass
 
-        return val_loss, 100*correct
+        return epoch_loss, epoch_metric
 
     @staticmethod
     def _add_pr_curve_tb(pred: Tensor,
                          class_label: List[Tensor],
                          class_probs: List[List[Tensor]],
                          writer: SummaryWriter,
-                         writer_tag: str="") -> None:
+                         writer_tag: str = "") -> None:
         """private function to add the PR curve
         to tensorboard"""
         probs = torch.cat([torch.stack(batch) for batch in
-                          class_probs]).cpu()
+                           class_probs]).cpu()
         labels = torch.cat(class_label).cpu()
         for class_index in range(len(pred[0])):
-            tensorboard_truth = 1*(labels == class_index).flatten()
+            tensorboard_truth = 1 * (labels == class_index).flatten()
             tensorboard_probs = probs[:, class_index]
-            #print(tensorboard_truth)
-            #print(tensorboard_probs)
+            # print(tensorboard_truth)
+            # print(tensorboard_probs)
             try:
-                writer.add_pr_curve(writer_tag+"/class = "+str(class_index),
-                                         tensorboard_truth,
-                                         tensorboard_probs,
-                                         global_step=0)
+                writer.add_pr_curve(writer_tag + "/class = " + str(class_index),
+                                    tensorboard_truth,
+                                    tensorboard_probs,
+                                    global_step=0)
             except AttributeError:
                 warnings.warn("Cannot store data in the PR curve")
             except ValueError:
                 warnings.warn("Cannot store data in the PR curve")
 
     @_add_data_to_tb
-    def _inner_loop(self,*,
+    def _inner_loop(self, *,
                     dl: DataLoader[Tuple[Tensor, Tensor]],
                     class_probs: List[List[Tensor]],
                     class_label: List[Tensor],
-                    loss: float,
-                    correct: float,
-                    writer: SummaryWriter,
-                    writer_tag: str) -> Tuple[List[Tensor], float, float]:
+                    writer_tag: str  # noqa
+                    ) -> Tuple[List[Tensor], float, float]:
         """private function used inside the test
         and validation loops"""
         pred_list = []
+        batch_metric_list = []
+        loss = 0.
         with torch.no_grad():
             for X, y in dl:
                 X = X.to(self.DEVICE)
                 y = y.to(self.DEVICE)
                 pred = self.model(X)
                 pred_list.append(pred)
-                class_probs_batch = [F.softmax(el, dim=0)
+                class_probs_batch = [f.softmax(el, dim=0)
                                      for el in pred]
                 class_probs.append(class_probs_batch)
                 loss += self.loss_fn(pred, y).item()
-                try:
-                    correct += (pred.argmax(1) == y).to(torch.float).sum().item()
-                except RuntimeError:
-                    correct += (pred.argmax(2) == y).to(torch.float).sum().item()
+                batch_metric = self.training_metric(pred, y)
+                batch_metric_list.append(batch_metric)
                 class_label.append(y)
-
-        return pred_list, loss, correct
+        epoch_metric = sum(batch_metric_list)/len(batch_metric_list)
+        epoch_loss = loss/len(batch_metric_list)
+        return pred_list, epoch_loss, epoch_metric
 
     def _init_profiler(self,
-                       profiling:bool,
-                       cross_validation:bool,
-                       n_epochs:int,
-                       k_folds:int) -> None:
+                       profiling: bool,
+                       cross_validation: bool,
+                       n_epochs: int,
+                       k_folds: int) -> None:
         """initialise the profler for profiling"""
         # profiling
+        active: int = 10
         if not cross_validation:
-            active = n_epochs-2
+            active = n_epochs - 2
         else:
-            active = k_folds*(n_epochs-2)
+            active = k_folds * (n_epochs - 2)
 
-        active:int = 10
         if profiling:
             try:
                 self.prof = torch.profiler.profile(
-                            activities=[
-                                torch.profiler.ProfilerActivity.CPU,
-                                torch.profiler.ProfilerActivity.CUDA],
-                            schedule=torch.profiler.schedule(wait=1, warmup=1, active=active, repeat=1),
-                            on_trace_ready=torch.profiler.tensorboard_trace_handler(
-                                os.path.join('.','runs', (self.model.__class__.__name__ +
-                                                          str(datetime.today())).replace(":","-")),
-                                worker_name='worker'),
-                            record_shapes=True,
-                            profile_memory=True
-                            #with_stack=True
+                    activities=[
+                        torch.profiler.ProfilerActivity.CPU,
+                        torch.profiler.ProfilerActivity.CUDA],
+                    schedule=torch.profiler.schedule(wait=1, warmup=1, active=active, repeat=1),
+                    on_trace_ready=torch.profiler.tensorboard_trace_handler(
+                        os.path.join('.', 'runs', (self.model.__class__.__name__ +
+                                                   str(datetime.today())).replace(":", "-")),
+                        worker_name='worker'),
+                    record_shapes=True,
+                    profile_memory=True
+                    # with_stack=True
                 )
             except AssertionError:
                 pass
-    
+
     def _init_optimizer_and_scheduler(self,
-                                      keep_training:bool,
-                                      cross_validation:bool,
-                                      optimizer:Type[Optimizer],
-                                      optimizers_param:Dict[str, Any],
-                                      lr_scheduler:Optional[Type[_LRScheduler]]=None,
-                                      scheduler_params:Optional[Dict[str, Any]]=None) -> None:
+                                      keep_training: bool,
+                                      cross_validation: bool,
+                                      optimizer: Type[Optimizer],
+                                      optimizers_param: Dict[str, Any],
+                                      lr_scheduler: Optional[Type[_LRScheduler]] = None,
+                                      scheduler_params: Optional[Dict[str, Any]] = None) -> None:
         """Reset or maintain the LR scheduler and the 
         optimizer depending on the training"""
         if not (self.check_has_trained and keep_training):
@@ -485,22 +475,22 @@ def _init_optimizer_and_scheduler(self,
             self.optimizer.__init__(self.model.parameters(), **dict_param)
             if lr_scheduler is not None:  # reset scheduler
                 self.scheduler = lr_scheduler(self.optimizer, **scheduler_params)
-    
+
     def train(self,
               optimizer: Type[Optimizer],
-              n_epochs: int=10,
-              cross_validation: bool=False,
-              optimizers_param: Optional[Dict[str, Any]]=None,
-              dataloaders_param: Optional[Dict[str, Any]]=None,
-              lr_scheduler: Optional[Type[_LRScheduler]]=None,
-              scheduler_params: Optional[Dict[str, Any]]=None,
-              optuna_params: Optional[Tuple[BaseTrial, str]]=None,
-              profiling: bool=False,
-              parallel_tpu: bool=False,
-              keep_training: bool=False,
-              store_grad_layer_hist: bool=False,
-              n_accumulated_grads: int=0,
-              writer_tag: str="") -> Tuple[float, float]:
+              n_epochs: int = 10,
+              cross_validation: bool = False,
+              optimizers_param: Optional[Dict[str, Any]] = None,
+              dataloaders_param: Optional[Dict[str, Any]] = None,
+              lr_scheduler: Optional[Type[_LRScheduler]] = None,
+              scheduler_params: Optional[Dict[str, Any]] = None,
+              optuna_params: Optional[Tuple[BaseTrial, str]] = None,
+              profiling: bool = False,
+              parallel_tpu: bool = False,
+              keep_training: bool = False,
+              store_grad_layer_hist: bool = False,
+              n_accumulated_grads: int = 0,
+              writer_tag: str = "") -> Tuple[float, float]:
         """Function to run all the training cycles.
 
         Args:
@@ -518,9 +508,6 @@ def train(self,
             optimizers_param (dict):
                 dictionary of the optimizers
                 parameters, e.g. `{"lr": 0.001}`
-            models_param (dict):
-                dictionary of the model
-                parameters
             lr_scheduler (torch.optim):
                 a learning rate scheduler class
             scheduler_params (dict):
@@ -566,7 +553,7 @@ def train(self,
         dl_tr = self.dataloaders[0]
         if optimizers_param is None:
             optimizers_param = {"lr": 0.001}
-            
+
         # dataloaders_param initialisation
         if dataloaders_param is None:
             if self.dataloaders[1] is not None:
@@ -577,7 +564,7 @@ def train(self,
         else:
             dataloaders_param_val = dataloaders_param.copy()
             dataloaders_param_tr = dataloaders_param.copy()
-        
+
         # scheduler_params initialisation
         if scheduler_params is None:
             scheduler_params = {}
@@ -585,21 +572,21 @@ def train(self,
         # LR scheduler
         if not (self.check_has_trained and keep_training):
             self.scheduler = None
-        
+
         # optuna gridsearch
         search_metric = None
         trial = None
-        if not optuna_params is None:
+        if optuna_params is not None:
             check_optuna = True
             trial, search_metric = optuna_params
         else:
             check_optuna = False
 
         # profiling
-        self._init_profiler(profiling, 
-                            cross_validation, 
-                            n_epochs, self.KFold_class.n_splits)
-        
+        self._init_profiler(profiling,
+                            cross_validation,
+                            n_epochs, self.k_fold_class.n_splits)
+
         # remove sampler to avoid conflicts with indexing
         # we will re-introduce the sampler when creating the indexing list
         try:
@@ -619,18 +606,18 @@ def train(self,
                 val_idx = self.dataloaders[1].sampler.indices  # type: ignore
             except AttributeError:
                 val_idx = list(range(len(self.dataloaders[1].dataset)))
-            #print(val_idx)
+            # print(val_idx)
             dl_val = torch.utils.data.DataLoader(self.dataloaders[1].dataset,
-                                                 #pin_memory=True,
+                                                 # pin_memory=True,
                                                  **dataloaders_param_val,
                                                  sampler=SubsetRandomSampler(val_idx))
             try:
                 tr_idx = self.dataloaders[0].sampler.indices  # type: ignore
             except AttributeError:
                 tr_idx = list(range(len(self.dataloaders[0].dataset)))
-            #print(tr_idx)
+            # print(tr_idx)
             dl_tr = torch.utils.data.DataLoader(self.dataloaders[0].dataset,
-                                                #pin_memory=True,
+                                                # pin_memory=True,
                                                 **dataloaders_param_tr,
                                                 sampler=SubsetRandomSampler(tr_idx))
         else:
@@ -638,21 +625,20 @@ def train(self,
                 data_idx = self.dataloaders[0].sampler.indices  # type: ignore
             except AttributeError:
                 data_idx = list(range(len(self.dataloaders[0].dataset)))
-            #print(data_idx)
+            # print(data_idx)
             tr_idx, val_idx = train_test_split(data_idx, test_size=0.2)
             dl_val = torch.utils.data.DataLoader(self.dataloaders[0].dataset,
-                                                 #pin_memory=True,
+                                                 # pin_memory=True,
                                                  **dataloaders_param_val,
                                                  sampler=SubsetRandomSampler(val_idx))
             dl_tr = torch.utils.data.DataLoader(self.dataloaders[0].dataset,
-                                                #pin_memory=True,
+                                                # pin_memory=True,
                                                 **dataloaders_param_tr,
                                                 sampler=SubsetRandomSampler(tr_idx))
 
         if cross_validation:
             mean_val_loss = []
             mean_val_acc = []
-            valloss, valacc = 0., 0.
             try:
                 data_idx = self.dataloaders[0].sampler.indices
                 labels_for_split = [self.dataloaders[0].dataset[i][-1] for i in data_idx]
@@ -660,54 +646,54 @@ def train(self,
                 data_idx = list(range(len(self.dataloaders[0].dataset)))
                 labels_for_split = [self.dataloaders[0].dataset[i][-1] for i in data_idx]
 
-            for fold, (tr_idx, val_idx) in enumerate(self.KFold_class.split(data_idx, labels_for_split)):
+            for fold, (tr_idx, val_idx) in enumerate(self.k_fold_class.split(data_idx, labels_for_split)):
                 # prints for class balance
-                #lab_tr_fold = [self.dataloaders[0].dataset[i][-1] for i in tr_idx]
-                #lab_val_fold = [self.dataloaders[0].dataset[i][-1] for i in val_idx]
-                #print("train labels:",[(i, lab_tr_fold.count(i)) for i in np.unique(np.array(lab_tr_fold))])
-                #print("val labels:",[(i, lab_val_fold.count(i)) for i in np.unique(np.array(lab_val_fold))])
-                #print("lenghts: ", len(lab_tr_fold), len(lab_val_fold))
+                # lab_tr_fold = [self.dataloaders[0].dataset[i][-1] for i in tr_idx]
+                # lab_val_fold = [self.dataloaders[0].dataset[i][-1] for i in val_idx]
+                # print("train labels:",[(i, lab_tr_fold.count(i)) for i in np.unique(np.array(lab_tr_fold))])
+                # print("val labels:",[(i, lab_val_fold.count(i)) for i in np.unique(np.array(lab_val_fold))])
+                # print("lenghts: ", len(lab_tr_fold), len(lab_val_fold))
                 self._init_optimizer_and_scheduler(keep_training, cross_validation,
                                                    optimizer, optimizers_param,
                                                    lr_scheduler, scheduler_params)
-                
+
                 # re-initialise data loaders
                 if len(self.dataloaders) == 3:
                     warnings.warn("Validation set is ignored in automatic Cross Validation")
                 dl_tr = torch.utils.data.DataLoader(self.dataloaders[0].dataset,
-                                                    #pin_memory=True,
+                                                    # pin_memory=True,
                                                     **dataloaders_param_tr,
                                                     sampler=SubsetRandomSampler(tr_idx))
                 dl_val = torch.utils.data.DataLoader(self.dataloaders[0].dataset,
-                                                     #pin_memory=True,
+                                                     # pin_memory=True,
                                                      **dataloaders_param_val,
                                                      sampler=SubsetRandomSampler(val_idx))
                 # print n-th fold
-                print("\n\n********** Fold ", fold+1, "**************")
+                print("\n\n********** Fold ", fold + 1, "**************")
                 # the training and validation loop
                 if parallel_tpu == False:
                     valloss, valacc = self._training_loops(n_epochs, dl_tr,
                                                            dl_val, lr_scheduler, self.scheduler,
                                                            check_optuna, search_metric,
                                                            trial, cross_validation,
-                                                           writer_tag + "/fold = " + str(fold+1))
+                                                           writer_tag + "/fold = " + str(fold + 1))
                 else:
                     valloss, valacc = self.parallel_tpu_training_loops(n_epochs, dl_tr,
-                                                       dl_val, optimizers_param, lr_scheduler,
-                                                       self.scheduler,
-                                                       check_optuna, search_metric,
-                                                       trial, cross_validation)
+                                                                       dl_val, optimizers_param, lr_scheduler,
+                                                                       self.scheduler,
+                                                                       check_optuna, search_metric,
+                                                                       trial, cross_validation)
 
                 mean_val_loss.append(valloss)
                 mean_val_acc.append(valacc)
             # mean of the validation and loss accuracies over folds
             if self.best_not_last:
                 valloss = min(np.array(_inner_refactor_scalars(self.val_loss_list_hparam,
-                                                  True,
-                                                  self.KFold_class.n_splits))[:,0])
+                                                               True,
+                                                               self.k_fold_class.n_splits))[:, 0])
                 valacc = max(np.array(_inner_refactor_scalars(self.val_acc_list_hparam,
-                                                 True,
-                                                 self.KFold_class.n_splits))[:,0])
+                                                              True,
+                                                              self.k_fold_class.n_splits))[:, 0])
             else:
                 valloss = np.mean(mean_val_loss)
                 valacc = np.mean(mean_val_acc)
@@ -717,17 +703,17 @@ def train(self,
                                                optimizer, optimizers_param,
                                                lr_scheduler, scheduler_params)
 
-            if parallel_tpu == False:
+            if not parallel_tpu:
                 valloss, valacc = self._training_loops(n_epochs, dl_tr,
-                                                   dl_val, lr_scheduler, self.scheduler,
-                                                   check_optuna, search_metric,
-                                                   trial, False, writer_tag)
+                                                       dl_val, lr_scheduler, self.scheduler,
+                                                       check_optuna, search_metric,
+                                                       trial, False, writer_tag)
             else:
                 valloss, valacc = self.parallel_tpu_training_loops(n_epochs, dl_tr,
-                                                   dl_val, optimizers_param, lr_scheduler,
-                                                   self.scheduler,
-                                                   check_optuna, search_metric,
-                                                   trial, False)
+                                                                   dl_val, optimizers_param, lr_scheduler,
+                                                                   self.scheduler,
+                                                                   check_optuna, search_metric,
+                                                                   trial, False)
 
         try:
             self.writer.flush()
@@ -738,18 +724,18 @@ def train(self,
 
         # put the mean of the cross_val
         return valloss, valacc
-        
+
     def _training_loops(self,
-                        n_epochs:int,
-                        dl_tr:DataLoader[Tuple[Tensor, Tensor]],
-                        dl_val:DataLoader[Tuple[Tensor, Tensor]],
-                        lr_scheduler: Optional[Type[_LRScheduler]]=None,
-                        scheduler: Optional[_LRScheduler]=None,
-                        check_optuna:bool=False,
-                        search_metric:Optional[str]=None,
-                        trial:Optional[BaseTrial]=None,
-                        cross_validation:bool=False,
-                        writer_tag:str="") -> Tuple[float, float]:
+                        n_epochs: int,
+                        dl_tr: DataLoader[Tuple[Tensor, Tensor]],
+                        dl_val: DataLoader[Tuple[Tensor, Tensor]],
+                        lr_scheduler: Optional[Type[_LRScheduler]] = None,
+                        scheduler: Optional[_LRScheduler] = None,
+                        check_optuna: bool = False,
+                        search_metric: Optional[str] = None,
+                        trial: Optional[BaseTrial] = None,
+                        cross_validation: bool = False,
+                        writer_tag: str = "") -> Tuple[float, float]:
         """private method to run the trainign loops
         
         Args:
@@ -760,12 +746,6 @@ def _training_loops(self,
             dl_val (torch.DataLoader):
                 validation dataloader
                 parameters, e.g. `{'batch_size': 32}`
-            optimizers_param (dict):
-                dictionary of the optimizers
-                parameters, e.g. `{"lr": 0.001}`
-            models_param (dict):
-                dictionary of the model
-                parameters
             lr_scheduler (torch.optim):
                 a learning rate scheduler class
             scheduler (torch.optim):
@@ -788,9 +768,9 @@ def _training_loops(self,
                 the validation loss and validation accuracy
         """
 
-        min_valloss, max_valacc = np.inf, 0.
+        valloss, valacc = np.inf, 0.
         for t in range(n_epochs):
-            print(f"Epoch {t+1}\n-------------------------------")
+            print(f"Epoch {t + 1}\n-------------------------------")
             self.val_epoch = t
             self.train_epoch = t
             self._train_loop(dl_tr, writer_tag)
@@ -800,18 +780,18 @@ def _training_loops(self,
                     lista = me.get_layers_param()
                     for k, item in lista.items():
                         try:
-                            self.writer.add_histogram(writer_tag + "/weights&biases/param/train/"+k,item,t)
-                            self.writer.add_histogram(writer_tag + "/weights&biases/param/train/log/"+k,
-                                                      torch.log(torch.abs(item)+1e-8),t)
+                            self.writer.add_histogram(writer_tag + "/weights&biases/param/train/" + k, item, t)
+                            self.writer.add_histogram(writer_tag + "/weights&biases/param/train/log/" + k,
+                                                      torch.log(torch.abs(item) + 1e-8), t)
                         except ValueError:
                             warnings.warn(f"The histogram is empty, most likely because your loss"
                                           f" is exploding. Try use gradient clipping.")
                     lista_grad = me.get_layers_grads()
-                    for k, item in zip(lista.keys(),lista_grad):
+                    for k, item in zip(lista.keys(), lista_grad):
                         try:
-                            self.writer.add_histogram(writer_tag + "/weights&biases/grads/train/"+k,item,t)
-                            self.writer.add_histogram(writer_tag + "/weights&biases/param/train/log/"+k,
-                                                      torch.log(torch.abs(item)+1e-8),t)
+                            self.writer.add_histogram(writer_tag + "/weights&biases/grads/train/" + k, item, t)
+                            self.writer.add_histogram(writer_tag + "/weights&biases/param/train/log/" + k,
+                                                      torch.log(torch.abs(item) + 1e-8), t)
                         except ValueError:
                             warnings.warn(f"The histogram is empty, most likely because your loss"
                                           f" is exploding. Try use gradient clipping.")
@@ -819,12 +799,12 @@ def _training_loops(self,
 
                 except AttributeError:
                     pass
-            self._run_pipe_hook(t+1, self.optimizer, me, self.writer)
+            self._run_pipe_hook(t + 1, self.optimizer, me, self.writer)
             valloss, valacc = self._val_loop(dl_val, writer_tag)
             # absolute best of loss and accuracy
             self.best_val_acc = max(self.best_val_acc, valacc)
             self.best_val_loss = min(self.best_val_loss, valloss)
-            #print(self.optimizer.param_groups[0]["lr"])
+            # print(self.optimizer.param_groups[0]["lr"])
             if not lr_scheduler is None:
                 scheduler.step()  # type: ignore
             # pruning trials
@@ -838,44 +818,45 @@ def _training_loops(self,
                     raise optuna.exceptions.TrialPruned()
         return valloss, valacc
 
-
     def parallel_tpu_training_loops(self,
-                                    n_epochs:int,
-                                    dl_tr_old:DataLoader[Tuple[Tensor, Tensor]],
-                                    dl_val_old:DataLoader[Tuple[Tensor, Tensor]],
+                                    n_epochs: int,
+                                    dl_tr_old: DataLoader[Tuple[Tensor, Tensor]],
+                                    dl_val_old: DataLoader[Tuple[Tensor, Tensor]],
                                     optimizers_param: Dict[str, Any],
-                                    lr_scheduler: Optional[Type[_LRScheduler]]=None,
-                                    scheduler: Optional[_LRScheduler]=None,
-                                    check_optuna: bool=False,
-                                    search_metric: Optional[str]=None,
-                                    trial:Optional[BaseTrial]=None,
-                                    cross_validation: bool=False) -> Tuple[float, float]:
+                                    lr_scheduler: Optional[Type[_LRScheduler]] = None,
+                                    scheduler: Optional[_LRScheduler] = None,
+                                    check_optuna: bool = False,
+                                    search_metric: Optional[str] = None,
+                                    trial: Optional[BaseTrial] = None,
+                                    cross_validation: bool = False) -> Tuple[float, float]:
         """Experimental function to run all the training cycles
         on colab TPUs in parallel.
         Note: ``cross_validation`` parameter as well as
         ``profiling`` are ignored.
 
         Args:
-            n_epochs (int):
+            n_epochs:
                 number of training epochs
-            dl_tr (torch.DataLoader):
+            dl_tr_old:
                 training dataloader
-            dl_val (torch.DataLoader):
+            dl_val_old:
                 validation dataloader
                 parameters, e.g. `{'batch_size': 32}`
-            lr_scheduler (torch.optim):
+            optimizers_param:
+                dictionary of parameters for the optimizers
+            lr_scheduler :
                 a learning rate scheduler class
-            scheduler (torch.optim):
+            scheduler:
                 the actual scheduler instance
-            check_optuna (bool):
+            check_optuna :
                 boolean to store the optuna results of
                 the trial
-            search_metric (string):
+            search_metric:
                 either ``'loss'`` or ``'accuracy'``, this
                 corresponds to the gridsearch criterion
-            trial (optuna.trial):
+            trial:
                 the optuna trial
-            cross_validation (bool, default False)
+            cross_validation :
                 the boolean flag for cross validation
 
         Returns:
@@ -885,8 +866,9 @@ def parallel_tpu_training_loops(self,
         self.val_loss = 0
         self.val_acc = 0
         warnings.warn("The tensorboard writer is ignored " +
-                      "for multi TPU processing. Also SAM optimisation"+
+                      "for multi TPU processing. Also SAM optimisation" +
                       " does not work for multi TPU training.")
+
         def map_fun_custom(index, flags):
             """map function for multi-processing"""
             device = xm.xla_device()
@@ -896,7 +878,7 @@ def map_fun_custom(index, flags):
 
             # initialize optimizer
             optimizer_class = type(self.optimizer)
-            optimizer = optimizer_class(model2.parameters(),**optimizers_param)
+            optimizer = optimizer_class(model2.parameters(), **optimizers_param)
 
             # define training and validation
             # distributed samplers and update
@@ -906,7 +888,7 @@ def map_fun_custom(index, flags):
                 num_replicas=xm.xrt_world_size(),
                 rank=xm.get_ordinal(),
                 shuffle=True)
-            
+
             dl_tr = torch.utils.data.DataLoader(
                 dl_tr_old.dataset,
                 num_workers=dl_tr_old.num_workers,
@@ -914,13 +896,13 @@ def map_fun_custom(index, flags):
                 sampler=train_sampler,
                 drop_last=dl_tr_old.drop_last
             )
-          
+
             val_sampler = torch.utils.data.distributed.DistributedSampler(
                 dl_val_old.dataset,
                 num_replicas=xm.xrt_world_size(),
                 rank=xm.get_ordinal(),
                 shuffle=False)
-            
+
             dl_val = torch.utils.data.DataLoader(
                 dl_val_old.dataset,
                 num_workers=dl_val_old.num_workers,
@@ -929,20 +911,17 @@ def map_fun_custom(index, flags):
                 drop_last=dl_val_old.drop_last
             )
 
-
-            #train loop
+            # train loop
             for t in range(n_epochs):
                 model2.train()
                 para_train_loader = pl.ParallelLoader(dl_tr,
                                                       [device]).per_device_loader(device)
-                print(f"Epoch {t+1}\n-------------------------------")
+                print(f"Epoch {t + 1}\n-------------------------------")
                 self.val_epoch = t
                 self.train_epoch = t
 
                 # train batch loop
-                size = len(dl_tr.dataset)
-                steps = len(dl_tr)
-                loss = 0.    # arbitrary starting value to avoid nan loss
+                loss = 0.
                 correct = 0.
                 tik = time.time()
                 # for batch, (X, y) in enumerate(self.dataloaders[0]):
@@ -955,7 +934,7 @@ def map_fun_custom(index, flags):
                         correct += (pred.argmax(2) == y).to(torch.float).sum().item()
                     loss = self.loss_fn(pred, y)
                     # Save to tensorboard
-                    #self.writer.add_scalar("Parallel" + "/Loss/train",
+                    # self.writer.add_scalar("Parallel" + "/Loss/train",
                     #                       loss.cpu(),
                     #                       self.train_epoch*dl_tr.batch_size + batch)
                     # Backpropagation
@@ -964,19 +943,17 @@ def map_fun_custom(index, flags):
 
                     xm.optimizer_step(self.optimizer)
 
-                #self.writer.flush()
                 # train accuracy:
-                correct /= len(dl_tr)*dl_tr.batch_size
-                print("Train accuracy at epoch ",t," : ", correct)
-                print(f"\nTime taken for this epoch: {round(time.time()-tik):.2f}s")
+                correct /= len(dl_tr) * dl_tr.batch_size
+                print("Train accuracy at epoch ", t, " : ", correct)
+                print(f"\nTime taken for this epoch: {round(time.time() - tik):.2f}s")
 
                 # evaluation
                 model2.eval()
-                
-                size = len(dl_val.dataset)
+
                 loss, correct = 0., 0.
-                class_label:List[Tensor] = []
-                class_probs:List[List[Tensor]] = []
+                class_label: List[Tensor] = []
+                class_probs: List[List[Tensor]] = []
 
                 pred = 0.
                 para_valid_loader = pl.ParallelLoader(dl_val,
@@ -985,7 +962,7 @@ def map_fun_custom(index, flags):
                     # per batch!!
                     for X, y in para_valid_loader:
                         pred = model2(X)
-                        class_probs_batch = [F.softmax(el, dim=0)
+                        class_probs_batch = [f.softmax(el, dim=0)
                                              for el in pred]
                         class_probs.append(class_probs_batch)
                         loss += self.loss_fn(pred, y).item()
@@ -995,17 +972,17 @@ def map_fun_custom(index, flags):
                             correct += (pred.argmax(2) == y).to(torch.float).sum().item()
                         class_label.append(y)
                     # add data to tensorboard
-                    #self._add_pr_curve_tb(pred, class_label, class_probs, "validation")
+                    # self._add_pr_curve_tb(pred, class_label, class_probs, "validation")
 
-                #self.writer.add_scalar("Parallel " + "/Accuracy/validation", correct, self.val_epoch)
-                print(f"Validation results: \n Accuracy: {(100*correct):.2f}%, \
+                # self.writer.add_scalar("Parallel " + "/Accuracy/validation", correct, self.val_epoch)
+                print(f"Validation results: \n Accuracy: {(100 * correct):.2f}%, \
                         Avg loss: {loss:>8f} \n")
 
-                #self.writer.flush()
-              
-                if not lr_scheduler is None:
+                # self.writer.flush()
+
+                if lr_scheduler is not None:
                     scheduler.step()
-                if not self.prof is None:
+                if self.prof is not None:
                     self.prof.step()
 
                 if check_optuna and not cross_validation:
@@ -1017,17 +994,17 @@ def map_fun_custom(index, flags):
                     if trial.should_prune():
                         raise optuna.exceptions.TrialPruned()
             self.val_loss += loss
-            self.val_acc += correct*100
+            self.val_acc += correct * 100
 
-        flags:Dict[Any, Any] = {}
-        self.val_acc /= len(dl_val_old)*dl_val_old.batch_size
+        flags: Dict[Any, Any] = {}
+        self.val_acc /= len(dl_val_old) * dl_val_old.batch_size
         self.val_loss /= len(dl_val_old)
         xmp.spawn(map_fun_custom, args=(flags,), nprocs=8, start_method='fork')
         return self.val_loss, self.val_acc
 
     def evaluate_classification(self,
-                                num_class:int=None,
-                                dl:DataLoader[Tuple[Tensor, Tensor]]=None
+                                num_class: int = None,
+                                dl: DataLoader[Tuple[Tensor, Tensor]] = None
                                 ) -> Tuple[float, float, np.ndarray]:
         """Method to evaluate the performance of the model.
         
@@ -1044,8 +1021,8 @@ def evaluate_classification(self,
         """
         if dl is None:
             dl = self.dataloaders[0]
-        class_probs:List[List[Tensor]] = []
-        class_label:List[Tensor] = []
+        class_probs: List[List[Tensor]] = []
+        class_label: List[Tensor] = []
         loss = 0.
         correct = 0.
         confusion_matrix = np.zeros((num_class, num_class))
@@ -1055,7 +1032,7 @@ def evaluate_classification(self,
                 X = X.to(DEVICE)
                 y = y.to(DEVICE)
                 pred = self.model(X)
-                class_probs_batch = [F.softmax(el, dim=0)
+                class_probs_batch = [f.softmax(el, dim=0)
                                      for el in pred]
                 class_probs.append(class_probs_batch)
                 loss += self.loss_fn(pred, y).item()
@@ -1064,16 +1041,15 @@ def evaluate_classification(self,
                 class_label.append(y)
                 for t, p in zip(y.view(-1), pred.argmax(1).view(-1)):
                     confusion_matrix[t.long(), p.long()] += 1
-        correct /= len(dl)*dl.batch_size
+        correct /= len(dl) * dl.batch_size
         loss /= len(dl)
-        return 100*correct, loss, confusion_matrix
-
+        return 100 * correct, loss, confusion_matrix
 
     def register_pipe_hook(self,
-                           callable:Callable[[Optimizer, int,
-                                              ModelExtractor,
-                                              SummaryWriter],
-                                             Any]) -> None:
+                           callable: Callable[[Optimizer, int,
+                                               ModelExtractor,
+                                               SummaryWriter],
+                                              Any]) -> None:
         """This method registers a function that
         will be called after each trainign step.
 
@@ -1088,29 +1064,28 @@ def register_pipe_hook(self,
                 the function to register"""
         self.registered_hook = callable
 
-
     def _run_pipe_hook(self,
-                       epoch:int,
-                       optim:Optimizer,
-                       me:ModelExtractor,
+                       epoch: int,
+                       optim: Optimizer,
+                       me: ModelExtractor,
                        writer: SummaryWriter) -> None:
         """private method that runs the hooked
         function at every epoch, after the single training loop"""
         if self.registered_hook is not None:
             self.registered_hook(epoch, optim, me, writer)
-            
+
     @staticmethod
     def copy_dataloader_params(original_dataloader: DataLoader[Tuple[Tensor, Tensor]]) -> Dict[str, Any]:
         """returns the dict of init parameters"""
         return {"batch_size": original_dataloader.batch_size,
-                #"batch_sampler": original_dataloader.batch_sampler, 
-                "num_workers": original_dataloader.num_workers, 
-                "collate_fn": original_dataloader.collate_fn, 
-                "pin_memory": original_dataloader.pin_memory, 
-                "drop_last": original_dataloader.drop_last, 
-                "timeout": original_dataloader.timeout, 
-                "worker_init_fn": original_dataloader.worker_init_fn, 
-                "multiprocessing_context": original_dataloader.multiprocessing_context, 
-                "generator": original_dataloader.generator, 
-                "prefetch_factor": original_dataloader.prefetch_factor, 
+                # "batch_sampler": original_dataloader.batch_sampler,
+                "num_workers": original_dataloader.num_workers,
+                "collate_fn": original_dataloader.collate_fn,
+                "pin_memory": original_dataloader.pin_memory,
+                "drop_last": original_dataloader.drop_last,
+                "timeout": original_dataloader.timeout,
+                "worker_init_fn": original_dataloader.worker_init_fn,
+                "multiprocessing_context": original_dataloader.multiprocessing_context,
+                "generator": original_dataloader.generator,
+                "prefetch_factor": original_dataloader.prefetch_factor,
                 "persistent_workers": original_dataloader.persistent_workers}
diff --git a/gdeep/trainer/trainer_config.py b/gdeep/trainer/trainer_config.py
new file mode 100644
index 00000000..2306145b
--- /dev/null
+++ b/gdeep/trainer/trainer_config.py
@@ -0,0 +1,75 @@
+from dataclasses import dataclass
+from typing import Type, Optional, Dict, Any, Tuple
+
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler  # noqa
+from optuna.trial._base import BaseTrial  # noqa
+
+
+@dataclass
+class TrainerConfig:
+    """Configuration class that contains the
+    parameters of the Trainer and of the Benchmark class
+
+    Args:
+        optimizer:
+            optimizer class, not the instance
+        n_epochs:
+            number of training epochs
+        optimizers_params:
+            dictionary of the optimizers
+            parameters, e.g. `{"lr": 0.001}`
+        dataloaders_params:
+            dictionary of the dataloaders
+            parameters
+        models_hyperparams:
+            dictionary of the model
+            parameters
+        lr_scheduler:
+            a learning rate scheduler class
+        schedulers_params:
+            learning rate scheduler parameters
+        optuna_params:
+            a tuple with the optuna trial and the
+            search metric (a string).
+        profiling:
+            whether or not you want to activate the
+            profiler
+        parallel_tpu:
+            boolean value to run the computations
+            on multiple TPUs
+        keep_training:
+            This flag allows to restart a training from
+            the existing optimizer as well as the
+            existing model
+        store_grad_layer_hist:
+            This flag allows to store the gradients
+            and the layer values in tensorboard for
+            each epoch
+        n_accumulated_grads:
+            this is the number of accumated grads. It
+            is taken into account only for positive integers
+        writer_tag:
+            tag to prepend to the ouput
+            on tensorboard
+    """
+    optimizer: Type[Optimizer]
+    n_epochs: int = 10
+    cross_validation: bool = False
+    optimizers_param: Optional[Dict[str, Any]] = None
+    dataloaders_param: Optional[Dict[str, Any]] = None
+    lr_scheduler: Optional[Type[_LRScheduler]] = None
+    scheduler_params: Optional[Dict[str, Any]] = None
+    optuna_params: Optional[Tuple[BaseTrial, str]] = None
+    profiling: bool = False
+    parallel_tpu: bool = False
+    keep_training: bool = False
+    store_grad_layer_hist: bool = False
+    n_accumulated_grads: int = 0
+    writer_tag: str = ""
+
+    def to_dict(self) -> Dict[str, Any]:
+        """method to transform the config file into a dictionary"""
+        return {
+            k: v for k, v in self.__dict__.items() if not k.startswith("_")
+        }
diff --git a/gdeep/utility/optimisation/persistence_grad.py b/gdeep/utility/optimisation/persistence_grad.py
index 5206d4b0..11c1996b 100644
--- a/gdeep/utility/optimisation/persistence_grad.py
+++ b/gdeep/utility/optimisation/persistence_grad.py
@@ -1,64 +1,68 @@
+from functools import reduce
+from itertools import chain, combinations
+import multiprocessing
+from typing import Iterator, Any, Callable, Tuple, List, Optional
+
 import torch
 from torch import optim
-#from gtda.homology import VietorisRipsPersistence as vrp
-#from gtda.homology import FlagserPersistence as flp
 from gph.python import ripser_parallel
-from gtda.homology import FlagserPersistence as flp
+from gtda.homology import FlagserPersistence as Flp
 import plotly.figure_factory as ff
 import plotly.graph_objects as go
-from itertools import chain, combinations
+from plotly.graph_objects import Figure
 import numpy as np
 from tqdm import tqdm
-import multiprocessing
 import operator as op
-from functools import reduce
-from typing import Iterator
 
+from gdeep.utility import DEVICE
 
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-    print("Using GPU!")
-else:
-    DEVICE = torch.device("cpu")
+Tensor = torch.Tensor
+Array = np.ndarray
 
 
-class PersistenceGradient():
-    '''This class computes the gradient of the persistence
+class PersistenceGradient:
+    """This class computes the gradient of the persistence
     diagram with respect to a point cloud. The algorithms has
     first been developed in https://arxiv.org/abs/2010.08356 .
 
-    Discalimer: this algorithm works well for generic point clouds.
+    Disclaimer: this algorithm works well for generic point clouds.
     In case your point cloud has many simplices with same
     filtration values, the matching of the points to the persistent
     features may fail to disambiguate.
 
     Args:
-        zeta (float): 
+        zeta :
             the relative weight of the regularisation part
-            of the `persistence_function`
-        homology_dimensions (tuple): 
+            of the ``persistence_function``
+        homology_dimensions :
             tuple of homology dimensions
-        collapse_edges (bool, default: False): 
+        collapse_edges :
             whether to use Collapse or not. Not implemented yet.
-        max_edge_length (float or np.inf): 
+        max_edge_length :
             the maximum edge length
             to be consider not infinity
-        approx_digits (int): 
+        approx_digits :
             digits after which to trunc floats for
             list comparison
-        metric (string): either `"euclidean"` or `"precomputed"`. 
-            The second one is in case of X being 
+        metric :
+            either ``"euclidean"`` or ``"precomputed"``.
+            The second one is in case of ``x`` being
             the pairwise-distance matrix or
-            the adjaceny matrix of a graph.
-        directed (bool): whether the input graph is a directed graph
-            or not. Relevant only if `metric = "precomputed"`
-        
-    '''
+            the adjacency matrix of a graph.
+        directed :
+            whether the input graph is a directed graph
+            or not. Relevant only if ``metric = "precomputed"``
+
+    """
 
-    def __init__(self, zeta: float = 0.5, homology_dimensions: tuple = (0, 1),
-                 collapse_edges: bool = False, max_edge_length: float = np.inf,
-                 approx_digits: int = 6, metric: str = "euclidean",
-                 directed: bool = False):
+    def __init__(self,
+                 zeta: float = 0.5,
+                 homology_dimensions: List[int] = (0, 1),
+                 collapse_edges: bool = False,
+                 max_edge_length: float = np.inf,
+                 approx_digits: int = 6,
+                 metric: str = "euclidean",
+                 directed: bool = False) -> None:
 
         self.collapse_edges = collapse_edges
         self.max_edge_length = max_edge_length
@@ -67,19 +71,25 @@ def __init__(self, zeta: float = 0.5, homology_dimensions: tuple = (0, 1),
         self.approx_digits = approx_digits
         self.zeta = zeta
         self.homology_dimensions = homology_dimensions
+        self.dist_mat: Optional[Array] = None
 
     @staticmethod
-    def powerset(iterable, max_length: int) -> Iterator:
-        '''powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)
-        up to `max_length`.'''
+    def powerset(iterable: Iterator,
+                 max_length: int) -> Iterator:
+        """powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)
+        up to `max_length`."""
         s = list(iterable)
         return chain.from_iterable(combinations(s, r) for
                                    r in range(0, max_length + 1))
 
     @staticmethod
-    def _parallel_apply_along_axis(func1d, axis, arr, *args, **kwargs):
+    def _parallel_apply_along_axis(func1d: Callable,
+                                   axis: int,
+                                   arr: Array,
+                                   *args: Any,
+                                   **kwargs: Any) -> Array:
         """
-        Like numpy.apply_along_axis(), but takes advantage of multiple
+        Like ``numpy.apply_along_axis()``, but takes advantage of multiple
         cores.
         """
         # Effective axis where apply_along_axis() will be applied by each
@@ -97,16 +107,16 @@ def _parallel_apply_along_axis(func1d, axis, arr, *args, **kwargs):
         # Freeing the workers:
         pool.close()
         pool.join()
-
         return np.concatenate(individual_results)
 
-    def _simplicial_pairs_of_indices(self, X):
-        '''Private function to compute the pair of indices in X to
-        matching the simplices.'''
-        simplices = list(self.powerset(list(range(0, len(X))),
+    def _simplicial_pairs_of_indices(self, x: Tensor) -> Tensor:
+        """Private function to compute the pair of indices in X to
+        matching the simplices.
+        """
+        simplices = list(self.powerset(list(range(0, len(x))),
                          max(self.homology_dimensions) + 2))[1:]
         simplices_array = np.array(simplices, dtype=object).reshape(-1, 1)
-        comb_number = comb(max(self.homology_dimensions)+2, 2)
+        comb_number = comb(max(self.homology_dimensions) + 2, 2)
         # the current computation bottleneck
         if len(simplices_array) > 10000000:
             pairs_of_indices = self._parallel_apply_along_axis(
@@ -119,17 +129,18 @@ def _simplicial_pairs_of_indices(self, X):
                                                    comb_number)
         return torch.tensor(pairs_of_indices, dtype=torch.int64)
 
-    def phi(self, X: torch.Tensor) -> torch.Tensor:
-        '''This function is from $(R^d)^n$ to $R^{|K|}$,
+    def phi(self, x: Tensor) -> Tensor:
+        """This function is from :math:`(R^d)^n` to :math:`R^{|K|}`,
         where K is the top simplicial complex of the VR filtration.
         It is defined as:
-        $\Phi_{\sigma}(X)=max_{i,j \in \sigma}||x_i-x_j||.$ '''
+        :math:`\\Phi_{\\sigma}(X)=max_{i,j \\in \\sigma}||x_i-x_j||.`
+        """
 
         if self.metric == "precomputed":
-            self.dist_mat = X
+            self.dist_mat = x
         else:
-            self.dist_mat = torch.cdist(X, X)
-        simplicial_pairs = self._simplicial_pairs_of_indices(X).reshape(-1, 2)
+            self.dist_mat = torch.cdist(x, x)
+        simplicial_pairs = self._simplicial_pairs_of_indices(x).reshape(-1, 2)
         ks = simplicial_pairs[:, 0]
         js = simplicial_pairs[:, 1]
         comb_number = comb(max(self.homology_dimensions) + 2, 2)
@@ -141,55 +152,62 @@ def phi(self, X: torch.Tensor) -> torch.Tensor:
                                                            0, ks),
                           1, js.reshape(-1, 1))).reshape(-1,
                                                          comb_number), 1)[0]
-        lista = torch.sort(lista)[0]  # not inplace
+        lista = torch.sort(lista)[0]
         return lista
 
-    def _compute_pairs(self, X):
-        '''Use giotto-tda to compute homology (b,d) pairs
+    @staticmethod
+    def _compute_pairs(x: Tensor) -> Tuple[Tensor, Tensor]:
+        """Use giotto-tda to compute homology (b,d) pairs
 
         Args:
-            X (tensor): this is the input point cloud or the input
+            x:
+                this is the input point cloud or the input
                 pairwise distance or the adjacency matrix of a graph
 
         Returns:
-            pairs (tensor): this tensor is the tensor obtained from
-                the list of pairs (b,d)
-            homology_dim (tensor): this 1d tensor contains the homology
-                group index'''
+            Tensor, Tensor:
+                the first tensor is the tensor obtained from
+                the list of pairs ``(b,d)``; the second 1d tensor
+                contains the homology
+                group index"""
 
-        dgms = flp().fit_transform([X.detach().numpy()])
+        dgms = Flp().fit_transform([x.detach().numpy()])
         pairs = dgms[0]
         return pairs[:, :2], pairs[:, 2]
 
-    def _persistence(self, X):
-        '''This function computes the persistence permutation.
+    def _persistence(self, x: Tensor) -> Tuple[Array, Tensor, Tensor]:
+        """This function computes the persistence permutation.
         This permutation permutes the filtration values and matches
         them as follows:
-        $Pers:Filt_K \subset \mathbb R^{|K|} \to (\mathbb R^2)^p
-        \times \mathbb R^q, \Phi(X) \mapsto D = \cup_i^p
-        (\Phi_{\sigma_i1}(X) ,\Phi_{\sigma_i2}(X) )
-        \times \cup_j^q (\Phi_{\sigma_j}(X),+\infty).$
+        :math:`Pers:Filt_K \\subset \\mathbb R^{|K|} \\to (\\mathbb R^2)^p
+        \\times \\mathbb R^q, \\Phi(X) \\mapsto D = \\cup_i^p
+        (\\Phi_{\\sigma_i1}(X) ,\\Phi_{\\sigma_i2}(X) )
+        \\times \\cup_j^q (\\Phi_{\\sigma_j}(X),+\\infty).`
 
         Args:
-            X (np.array): this is the point cloud
+            x:
+                this is the point cloud
 
         Returns:
-            persistence_pairs (2darray): this is an array of pairs (i,j). `i`
-                and `j` are the indices in the list `phi` associated to a
-                positive-negative pair. When `j` is equal to `-1`, it means
+            persistence_pairs :
+                this is an array of pairs (i,j). ``i``
+                and ``j`` are the indices in the list ``phi`` associated to a
+                positive-negative pair. When ``j`` is equal to ``-1``, it means
                 that the match was not found and the feature
-                is infinitely persistent. When `(i,j)=(-1,-1)`, simply
+                is infinitely persistent. When ``(i,j)=(-1,-1)``, simply
                 discard the pair
-            phi (tensor): this is the 1d tensor of all filtration values
-            homology_dims (tensor): this 1d tensor contains the homology
+            phi :
+                this is the 1d tensor of all filtration values
+            homology_dims :
+                this 1d tensor contains the homology
                 group index
-        '''
-        phi = self.phi(X)
-        pairs, homology_dims = self._compute_pairs(X)
-        approx_pairs = list(np.around(pairs, self.approx_digits).flatten())
+        """
+        phi = self.phi(x)
+        pairs, homology_dims = self._compute_pairs(x)
+        approx_pairs: List[float] = list(np.around(pairs, self.approx_digits).flatten())
         num_approx = 10**self.approx_digits
         inv_num_approx = 10**(-self.approx_digits)
-        phi_approx = torch.round(phi*num_approx)*inv_num_approx
+        phi_approx: Tensor = torch.round(phi*num_approx)*inv_num_approx
         # find the indices of phi_approx elements that are in approx_pairs
         indices_in_phi_of_pairs = []
         for i in approx_pairs:
@@ -200,8 +218,8 @@ def _persistence(self, X):
                                            dtype=np.int32)
         for i in indices_in_phi_of_pairs:
             try:
-                temp_index = approx_pairs.index(round(phi[i].item(),
-                                                self.approx_digits))
+                temp_index: int = approx_pairs.index(round(phi[i].item(),
+                                                     self.approx_digits))
                 approx_pairs[temp_index] = float('inf')
                 persistence_pairs_array[temp_index//2,
                                         temp_index % 2] = int(i)
@@ -209,27 +227,27 @@ def _persistence(self, X):
                 pass
         return persistence_pairs_array, phi, homology_dims
     
-    def _computing_persistence_with_gph(self, X: torch.Tensor) -> list:
+    def _computing_persistence_with_gph(self, xx: Tensor) -> List:
         """This method accepts the pointcloud and returns the
         persistence diagram in the following form
-        $Pers:Filt_K \subset \mathbb R^{|K|} \to (\mathbb R^2)^p
-        \times \mathbb R^q, \Phi(X) \mapsto D = \cup_i^p
-        (\Phi_{\sigma_i1}(X) ,\Phi_{\sigma_i2}(X) )
-        \times \cup_j^q (\Phi_{\sigma_j}(X),+\infty).$
-        The persstence diagram ca be readily used for 
+        :math:`Pers:Filt_K \\subset \\mathbb R^{|K|} \\to (\\mathbb R^2)^p
+        \\times \\mathbb R^q, \\Phi(X) \\mapsto D = \\cup_i^p
+        (\\Phi_{\\sigma_i1}(X) ,\\Phi_{\\sigma_i2}(X) )
+        \\times \\cup_j^q (\\Phi_{\\sigma_j}(X),+\\infty).`
+        The persistence diagram can be readily used for
         gradient descent.
 
         Args:
-            X (torch.tensor):
-                point cloud
+            xx:
+                point cloud with ``shape = (n_points, n_features)``
 
         Returns:
             list of shape (n, 3):
-                Persistence pairs (correspondig to the
+                Persistence pairs (corresponding to the
                 first 2 dimensions) where the last dimension 
                 contains the homology dimension
         """
-        output = ripser_parallel(X.detach().numpy(),
+        output = ripser_parallel(xx.detach().numpy(),
                                  maxdim=max(self.homology_dimensions),
                                  thresh=self.max_edge_length,
                                  coeff=2,
@@ -239,61 +257,64 @@ def _computing_persistence_with_gph(self, X: torch.Tensor) -> list:
                                  return_generators=True)
         
         persistence_pairs = []
-        #print(output["gens"])
+
         for dim in self.homology_dimensions:
             if dim == 0:
-                persistence_pairs += [(0, torch.norm(X[x[1]]-X[x[2]]),
+                persistence_pairs += [(0, torch.norm(xx[x[1]]-xx[x[2]]),
                                       0) for x in output["gens"][dim]]
             else:
-                persistence_pairs += [(torch.norm(X[x[1]]-X[x[0]]), 
-                                      torch.norm(X[x[3]]-X[x[2]]), 
+                persistence_pairs += [(torch.norm(xx[x[1]]-xx[x[0]]),
+                                      torch.norm(xx[x[3]]-xx[x[2]]),
                                       dim) for x in output["gens"][1][dim-1]]
         return persistence_pairs
-        
 
-    def persistence_function(self, X: torch.Tensor) -> torch.Tensor:
-        '''This is the Loss functon to optimise.
-        $L=-\sum_i^p |\epsilon_{i2}-\epsilon_{i1}|+
-        \lambda \sum_{x in X} ||x||_2^2$
+    def persistence_function(self, xx: Tensor) -> Tensor:
+        """This is the Loss function to optimise.
+        :math:`L=-\\sum_i^p |\\epsilon_{i2}-\\epsilon_{i1}|+
+        \\lambda \\sum_{x in X} ||x||_2^2`
         It is composed of a regularisation term and a
         function on the filtration values that is (p,q)-permutation
-        invariant.'''
+        invariant."""
         out = 0
         # this is much slower
         if self.directed and self.metric == "precomputed":
-            persistence_array, phi, _ = self._persistence(X)
+            persistence_array, phi, _ = self._persistence(xx)
             for item in persistence_array:
                 if item[1] != -1:
                     out += (phi[item[1]]-phi[item[0]])
         else:
-            persistence_array = self._computing_persistence_with_gph(X)
+            persistence_array = self._computing_persistence_with_gph(xx)
             for item in persistence_array:
                 out += item[1]-item[0]
-        reg = (X**2).sum()  # regularisation term
+        reg = (xx**2).sum()  # regularisation term
         return -out + self.zeta*reg  # maximise persistence
 
-    def SGD(self, X: torch.Tensor, lr: float = 0.01, n_epochs: int = 5):
-        '''This function is the core function of this class and uses the
-        SGD method to move points around in ordder to optimise
-        `persistence_function`
+    def sgd(self, xx: Tensor,
+            lr: float = 0.01,
+            n_epochs: int = 5) -> Tuple[Figure, Figure, List[float]]:
+        """This function is the core function of this class and uses the
+        SGD method to move points around in order to optimise
+        ``persistence_function``
 
         Args:
-            Xp (torch.tensor): 2d tensor representing the point cloud,
-                the first dimension is `n_points` while the second
-                `n_features`
-            lr (float): learning rate for the SGD
-            n_epochs (int): the number of gradient iterations
+            xx:
+                2d tensor representing the point cloud,
+                the first dimension is ``n_points`` while the second
+                ``n_features``
+            lr:
+                learning rate for the SGD
+            n_epochs:
+                the number of gradient iterations
 
         Returns:
-           fig : plotly `quiver_plot`
-           fig3d : plotly `cone_plot`
-           loss_val (list): loss function values over the epochs'''
-
-        if not type(X) == torch.Tensor:
-            X = torch.tensor(X)
-        X.to(DEVICE)
-        X.requires_grad = True
-        grads = torch.zeros_like(X)
+           fig, fig3d, loss_val:
+               respectively the plotly ``quiver_plot``, plotly ``cone_plot``
+               ad the list of loss function values over the epochs"""
+
+        if not type(xx) == Tensor:
+            xx = torch.tensor(xx)
+        xx.to(DEVICE)
+        xx.requires_grad = True
         x = np.array([])
         z = np.array([])
         y = np.array([])
@@ -301,30 +322,30 @@ def SGD(self, X: torch.Tensor, lr: float = 0.01, n_epochs: int = 5):
         v = np.array([])
         w = np.array([])
         loss_val = []
-        optimizer = optim.Adam([X], lr=lr)
+        optimizer = optim.Adam([xx], lr=lr)
         for _ in tqdm(range(n_epochs)):
             optimizer.zero_grad()
-            loss = self.persistence_function(X)
+            loss = self.persistence_function(xx)
             loss_val.append(loss.item())
-            x = np.concatenate((x, X.detach().numpy()[:, 0]))
-            y = np.concatenate((y, X.detach().numpy()[:, 1]))
+            x = np.concatenate((x, xx.detach().numpy()[:, 0]))
+            y = np.concatenate((y, xx.detach().numpy()[:, 1]))
             loss.backward()  # compute gradients and store them in Xp.grad
-            grads = -X.grad.detach()
+            grads = -xx.grad.detach()
             u = np.concatenate((u, 1/grads.norm(2,
                                 1).mean()*grads.numpy()[:, 0]))
             v = np.concatenate((v, 1/grads.norm(2,
                                 1).mean()*grads.numpy()[:, 1]))
             try:
-                z = np.concatenate((z, X.detach().numpy()[:, 2]))
+                z = np.concatenate((z, xx.detach().numpy()[:, 2]))
                 w = np.concatenate((w,
                                     1/grads.norm(2) *
                                     grads.numpy()[:, 2]))
             except IndexError:
-                z = np.concatenate((z, 0*X.detach().numpy()[:, 1]))
+                z = np.concatenate((z, 0*xx.detach().numpy()[:, 1]))
                 w = np.concatenate((w, 0*grads.numpy()[:, 1]))
             optimizer.step()
         fig = ff.create_quiver(x, y, u, v)
-        fig3d = go.Figure(data=go.Cone(
+        fig3d = Figure(data=go.Cone(
             x=x,
             y=y,
             z=z,
@@ -338,7 +359,7 @@ def SGD(self, X: torch.Tensor, lr: float = 0.01, n_epochs: int = 5):
 
 
 # this function is outside the main class to use multiprocessing
-def unpacking_apply_along_axis(tupl: tuple) -> np.ndarray:
+def unpacking_apply_along_axis(tupl: List[Any]) -> Array:
     """
     Like numpy.apply_along_axis(), but with arguments in a tuple
     instead.
@@ -352,21 +373,23 @@ def unpacking_apply_along_axis(tupl: tuple) -> np.ndarray:
     return np.apply_along_axis(func1d, axis, arr, *args, **kwargs)
 
 
-# this function is outside the main class to use multiprocessing
-def _combinations_with_single(tupl, max_length):
-    '''Private function to compute combinations also for singletons
+def _combinations_with_single(tupl, max_length: int) -> Array:
+    """Private function to compute combinations also for singletons
     with repetition. The un-intuitive shape of the output is
     becaue of the padding and vectorized operation happening
     in the next steps.
 
     Args:
-        tupl (list): this is a list with only one element and such
+        tupl :
+            this is a list with only one element and such
             element is a tuple
-        max_length (int): this is the paddding length.
+        max_length :
+            this is the paddding length.
 
     Returns:
-        padded list -- with (0,0) -- of all the combinations of indices
-            within a given simplex.'''
+        ndarray :
+            padded list -- with (0,0) -- of all the combinations of indices
+            within a given simplex."""
     tupl = tupl[0]  # the slice dimension will always contain 1 element
     if len(tupl) == 1:
         list1 = list(((tupl[0], tupl[0]), )) + \
@@ -378,7 +401,7 @@ def _combinations_with_single(tupl, max_length):
 
 
 def comb(nnn: int, kkk: int) -> int:
-    '''this function computes n!/(k!*(n-k)!) efficiently'''
+    """this function computes :math:`\\frac{n!}{(k!*(n-k)!)}` efficiently"""
     kkk = min(kkk, nnn-kkk)
     numer = reduce(op.mul, range(nnn, nnn-kkk, -1), 1)
     denom = reduce(op.mul, range(1, kkk+1), 1)
diff --git a/gdeep/utility/utils.py b/gdeep/utility/utils.py
index 62a1b786..9fca12ce 100644
--- a/gdeep/utility/utils.py
+++ b/gdeep/utility/utils.py
@@ -1,7 +1,6 @@
-from ctypes import Union
-#import imp
 from IPython import get_ipython  # type: ignore
 import base64
+from typing import Dict, Optional, Type, List
 import os
 import time
 import hashlib
@@ -9,9 +8,12 @@
 
 import torch
 from torch import nn
+from torch.utils.data import DataLoader
+from torch.optim import Optimizer
 
 
-def _are_compatible(model_dict, dataloaders_dict):
+def _are_compatible(model_dict: Dict[str, nn.Module],
+                    dataloaders_dict: Dict[str, List[DataLoader]]) -> bool:
     """utility function to check the compatibility of a model
     with a set of dataloaders `(dl_tr, dl_val, dl_ts)`
     """
@@ -25,11 +27,12 @@ def _are_compatible(model_dict, dataloaders_dict):
     else:
         return True
 
-def save_model_and_optimizer(model,
-                             model_name: str=None,
-                             trial_id: str=None,
-                             optimizer=None,
-                             store_pickle=False):
+
+def save_model_and_optimizer(model: nn.Module,
+                             model_name: Optional[str] = None,
+                             trial_id: Optional[str] = None,
+                             optimizer: Optional[Optimizer] = None,
+                             store_pickle: bool = False):
     """Save the model and the optimizer state_dict
 
     Args:
@@ -63,7 +66,7 @@ def save_model_and_optimizer(model,
         else:
             torch.save(model.state_dict(),
                        os.path.join("state_dicts",
-                                    model.__class__.__name__+"-"+trial_id+".pth"))
+                                    model.__class__.__name__ + "-" + trial_id + ".pth"))
     else:
         if store_pickle:
             torch.save(model,
@@ -72,27 +75,31 @@ def save_model_and_optimizer(model,
         else:
             torch.save(model.state_dict(),
                        os.path.join("state_dicts",
-                                    model_name +"-"+trial_id+".pth"))
+                                    model_name + "-" + trial_id + ".pth"))
     if optimizer is not None:
         torch.save(optimizer.state_dict(),
                    os.path.join("state_dicts",
-                                str(optimizer).replace("\n","").replace("(","").replace(":","").replace(")","")
-                                +"-"+trial_id+".pth"))
+                                str(optimizer).replace("\n", "").replace("(", "").replace(":", "").replace(")", "")
+                                + "-" + trial_id + ".pth"))
 
 
-def ensemble_wrapper(clss):
+def ensemble_wrapper(clss: Type):
     """function to wrap the ensemble estimators
-    of the ``torchensable`` library.
+    of the ``torchensemble`` library.
 
     The only argument is the estimator class. Then
     you can initialise the output of this function
     as you would normally do for the original
-    ``torchensebl``` class
+    ``torchensemble`` class
 
     Args
-        clss (type):
+        clss:
             the class of the estimator, like
             ``VotingClassifier`` for example
+    Returns:
+        type:
+            the initialised ensemble estimator
+            class that is compatible with giotto-deep
     """
 
     class NewEnsembleEstimator(clss):
@@ -108,15 +115,16 @@ def _inner_refactor_scalars(list_, cross_validation, k_folds):
     per epoch"""
     out = []
     for t in range(len(list_)):
-        lis = [x[0] for x in list_ if x[1]==t]
+        lis = [x[0] for x in list_ if x[1] == t]
         value = sum(lis)
         if len(lis) > 0:
             if cross_validation:
-                out.append([value/k_folds , t])
+                out.append([value / k_folds, t])
             else:
                 out.append([value, t])
     return out
 
+
 def is_notebook() -> bool:
     """Check if the current environment is a notebook
     
@@ -127,14 +135,15 @@ def is_notebook() -> bool:
     try:
         shell = get_ipython().__class__.__name__
         if shell == 'ZMQInteractiveShell':
-            return True   # Jupyter notebook or qtconsole
+            return True  # Jupyter notebook or qtconsole
         elif shell == 'TerminalInteractiveShell':
             return False  # Terminal running IPython
         else:
             return False  # Other type (?)
     except NameError:
-        return False      # Probably standard Python interpreter
-    
+        return False  # Probably standard Python interpreter
+
+
 def autoreload_if_notebook() -> None:
     """Autoreload the modules if the environment is a notebook
     
@@ -144,7 +153,8 @@ def autoreload_if_notebook() -> None:
     from IPython import get_ipython  # type: ignore
     get_ipython().magic('load_ext autoreload')
     get_ipython().magic('autoreload 2')
-    
+
+
 def _file_as_bytes(file) -> bytes:
     """Returns a bytes object representing the file
 
@@ -158,13 +168,16 @@ def _file_as_bytes(file) -> bytes:
     """
     with open(file, 'rb') as f:
         return f.read()
-    
+
+
 def get_checksum(file: str, encoding: str = "hex"):
     """Returns the checksum of the file
 
     Args:
-        file (str):
+        file :
             Path to the file
+        encoding:
+            string encoding like ``'hex'``
             
     Raises:
         ValueError: if the file does not exist
@@ -183,7 +196,7 @@ def get_checksum(file: str, encoding: str = "hex"):
     elif encoding == "base64":
         return base64.b64encode(
             bytes.fromhex(hashlib.md5(_file_as_bytes(file)).hexdigest())
-            )
+        )
     else:
         raise ValueError("encoding must be either 'hex' or 'base64'")
 
@@ -192,12 +205,12 @@ class KnownWarningSilencer:
     """silence all warnings within this ``with``
     statement with this class"""
 
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
-    def __enter__(self):
+    def __enter__(self) -> 'KnownWarningSilencer':
         warnings.filterwarnings("ignore")
         return self
 
     def __exit__(self, type, value, traceback):
-        warnings.filterwarnings("default")
\ No newline at end of file
+        warnings.filterwarnings("default")
diff --git a/gdeep/visualisation/__init__.py b/gdeep/visualisation/__init__.py
index e2898bd0..46ab2964 100644
--- a/gdeep/visualisation/__init__.py
+++ b/gdeep/visualisation/__init__.py
@@ -1,6 +1,6 @@
 from .persistence_activations import persistence_diagrams_of_activations
 from .plot_compactification import Compactification
-from .utility import plotly2tensor, png2tensor
+from .utils import plotly2tensor, png2tensor
 from .visualiser import Visualiser
 
 
diff --git a/gdeep/visualisation/persistence_activations.py b/gdeep/visualisation/persistence_activations.py
index 721ed558..84040d4d 100644
--- a/gdeep/visualisation/persistence_activations.py
+++ b/gdeep/visualisation/persistence_activations.py
@@ -1,17 +1,20 @@
-import torch
+import numpy
 from gtda.homology import VietorisRipsPersistence, \
     WeakAlphaPersistence
 from gtda.graphs import KNeighborsGraph, \
     GraphGeodesicDistance
+from typing import List
 
+Array = numpy.ndarray
 
-def knn_distance_matrix(X, k=3):
+
+def knn_distance_matrix(x: List, k: int = 3):
 
     """Returns list of distance matrices according
     to the k-NN graph distance of the datasets X
 
     Args:
-        X (ndarray):
+        x:
             the point cloud
         k (int):
             the numbers of neighbors for the k-NN graph
@@ -19,7 +22,7 @@ def knn_distance_matrix(X, k=3):
     """
 
     kng = KNeighborsGraph(n_neighbors=k)
-    adj = kng.fit_transform(X)
+    adj = kng.fit_transform(x)
     return GraphGeodesicDistance(directed=False).fit_transform(adj)
 
 
@@ -54,11 +57,11 @@ def persistence_diagrams_of_activations(activations_list,
             of the different layers
     """
     for i, activ in enumerate(activations_list):
-        if len(activ.shape) > 2:  # in caso of non FF layers
+        if len(activ.shape) > 2:  # in case of non FF layers
             activations_list[i] = activ.view(activ.shape[0], -1)
 
     if k > 0 and mode == 'VR':
-        VR = VietorisRipsPersistence(homology_dimensions=homology_dimensions,
+        vr = VietorisRipsPersistence(homology_dimensions=homology_dimensions,
                                      metric='precomputed',
                                      reduced_homology=False,
                                      infinity_values=20,
@@ -67,11 +70,11 @@ def persistence_diagrams_of_activations(activations_list,
         # infinity_values set to avoid troubles with
         # multiple topological features living indefinitely
     elif k == -1 and mode == 'alpha':
-        VR = WeakAlphaPersistence(homology_dimensions=homology_dimensions,
+        vr = WeakAlphaPersistence(homology_dimensions=homology_dimensions,
                                   reduced_homology=False,
                                   max_edge_length=max_edge_length)
     else:
-        VR = VietorisRipsPersistence(homology_dimensions=homology_dimensions,
+        vr = VietorisRipsPersistence(homology_dimensions=homology_dimensions,
                                      reduced_homology=False,
                                      max_edge_length=max_edge_length,
                                      collapse_edges=True)
@@ -81,8 +84,8 @@ def persistence_diagrams_of_activations(activations_list,
             activations_list[i] = activ.cpu()
         dist_matrix = knn_distance_matrix(activations_list,
                                           k=k)
-        persistence_diagrams = VR.fit_transform(dist_matrix)
+        persistence_diagrams = vr.fit_transform(dist_matrix)
     else:
-        persistence_diagrams = VR.fit_transform(activations_list)
+        persistence_diagrams = vr.fit_transform(activations_list)
 
     return persistence_diagrams
diff --git a/gdeep/visualisation/plot_compactification.py b/gdeep/visualisation/plot_compactification.py
index 31e8087a..413e5594 100644
--- a/gdeep/visualisation/plot_compactification.py
+++ b/gdeep/visualisation/plot_compactification.py
@@ -1,3 +1,4 @@
+from typing import List, Tuple, Union
 # deep learning
 import torch
 import numpy as np
@@ -14,37 +15,40 @@
 # ML
 from sklearn.metrics import pairwise_distances
 
+Tensor = torch.Tensor
+Array = np.ndarray
 
-class Compactification():
-    '''Class to compactify the feature space and get
+
+class Compactification:
+    """Class to compactify the feature space and get
     better topological insights.
 
     Args:
-        precision (float, default = 0.4):
+        precision :
             A tolerance on how close to the decision
             boundary the points shall be.
-        n_features (int, default None):
+        n_features :
             number of input features
-        n_samples (int, default 1000):
+        n_samples:
             number of samples per chart
-        n_epochs (int, default 5000):
+        n_epochs:
             number of epochs needed to push points to the
             decison boundary
-        boundary_tuple (list):
+        boundary_tuple :
             list of pairs (left,right).
             This list defines the boundaries in each coordinates
-        neural_net (nn.Module):
+        neural_net :
             the trained network of which to compute
             the boundary
-    '''
+    """
 
-    def __init__(self, precision=0.4,
+    def __init__(self, precision: float = 0.4,
                  n_features: int = None,
                  n_samples: int = 1000,
                  epsilon: float = 0.05,
                  n_epochs: int = 5000,
-                 boundary_tuple = None,
-                 neural_net: nn.Module = None):
+                 boundary_tuple: List[Tuple[float, float]] = None,
+                 neural_net: nn.Module = None) -> None:
 
         if n_features is None:
             raise ValueError("Need to specify the number of features.")
@@ -55,7 +59,7 @@ def __init__(self, precision=0.4,
         else:
             self.neural_net = neural_net
         if boundary_tuple is None:
-            self.boundary_tuple = [(-1, 1) for ii in range(self.n_features)]
+            self.boundary_tuple = [(-1., 1.) for _ in range(self.n_features)]
         else:
             self.boundary_tuple = boundary_tuple
         self.precision = precision
@@ -67,7 +71,8 @@ def __init__(self, precision=0.4,
         self.label_final = None
         self.check_compute_charts = False
 
-    def transition_to_patch(self, sample_points_tensor, i):
+    @staticmethod
+    def _transition_to_patch(sample_points_tensor: Tensor, i: int) -> Union[Tensor, Array]:
         stacking_list = []
         if i == -1:
             points = sample_points_tensor
@@ -77,13 +82,13 @@ def transition_to_patch(self, sample_points_tensor, i):
                 if dim == i:
                     stacking_list.append(1/sample_points_tensor[:, dim])
                 else:
-                    stacking_list.append(sample_points_tensor[:, dim] / \
-                        sample_points_tensor[:, i])
-            if type(sample_points_tensor) == torch.Tensor:
+                    stacking_list.append(sample_points_tensor[:, dim] /
+                                         sample_points_tensor[:, i])
+            if type(sample_points_tensor) == Tensor:
                 return torch.stack(stacking_list, dim=1)
             return np.stack(stacking_list, axis=1)
 
-    def compute_charts(self):
+    def _compute_charts(self) -> List[Tensor]:
         # compute boundary
         self.patches = []
 
@@ -95,27 +100,25 @@ def compute_charts(self):
                                                   n_samples=self.n_samples)
             plot_points_tensor = torch.from_numpy(sample_points()).float()
             # move to i-th patch
-            sample_points_tensor = self.transition_to_patch(plot_points_tensor,
-                                                            i)
+            sample_points_tensor = self._transition_to_patch(plot_points_tensor, i)
 
             # Using new gradient flow implementation
             gf = GradientFlowDecisionBoundaryCalculator(model=self.neural_net,
-                                                        initial_points=
-                                                        sample_points_tensor,
+                                                        initial_points=sample_points_tensor,
                                                         optimizer=lambda params:
                                                         torch.optim.Adam(params))
             gf.step(number_of_steps=self.n_epochs)
             res = gf.get_filtered_decision_boundary(delta=self.precision).detach().cpu()
             # back to 0-th patch
-            plot_points_tensor = self.transition_to_patch(res, i)
+            plot_points_tensor = self._transition_to_patch(res, i)
             self.patches.append(plot_points_tensor)
         self.check_compute_charts = True
         return self.patches
 
-    def create_final_distance_matrix(self):
-        '''This is the main function to call without arguments'''
+    def create_final_distance_matrix(self) -> Tuple[Array, List]:
+        """This is the main function to call without arguments"""
         if self.check_compute_charts is False:
-            self.compute_charts()
+            self._compute_charts()
         self.list_of_pts_in_patches = []
         for i in range(len(self.patches)):
             list_of_pts_per_patch = []
@@ -123,8 +126,8 @@ def create_final_distance_matrix(self):
                 if j == i:
                     list_of_pts_per_patch.append(pts)
                 else:
-                    list_of_pts_per_patch.append(self.transition_to_patch(\
-                        self.transition_to_patch(pts, j-1), i-1))
+                    list_of_pts_per_patch.append(self._transition_to_patch(
+                        self._transition_to_patch(pts, j-1), i-1))
             self.list_of_pts_in_patches.append(np.concatenate(list_of_pts_per_patch))
 
         list_of_distances = []
@@ -138,15 +141,14 @@ def create_final_distance_matrix(self):
             self.label_final += list(j*np.ones(len(patch)))
         return d_final, self.label_final
 
-    def plot_chart(self, i):
-        '''This functions plots the points in each chart.
+    def plot_chart(self, i: int) -> None:
+        """This functions plots the points in each chart.
 
         Args:
-            i (int):
-                the chart index, from `-1` to `n_features`
-        '''
-        df_plot = pd.DataFrame(self.list_of_pts_in_patches[i], columns =
-                               ["x" + str(j) for j in
+            i :
+                the chart index, from ``-1`` to ``n_features``
+        """
+        df_plot = pd.DataFrame(self.list_of_pts_in_patches[i], columns=["x" + str(j) for j in
                                range(self.list_of_pts_in_patches[i].shape[1])])
         df_plot["label"] = self.label_final
         if self.list_of_pts_in_patches[i].shape[1] == 2:
diff --git a/gdeep/visualisation/utility.py b/gdeep/visualisation/utils.py
similarity index 100%
rename from gdeep/visualisation/utility.py
rename to gdeep/visualisation/utils.py
index 6498906e..df4244e7 100644
--- a/gdeep/visualisation/utility.py
+++ b/gdeep/visualisation/utils.py
@@ -1,7 +1,7 @@
+import os
 
 from PIL import Image
 import numpy as np
-import os
 from plotly.io import write_image
 import torch
 
diff --git a/gdeep/visualisation/visualiser.py b/gdeep/visualisation/visualiser.py
index b5faecfd..280d1246 100644
--- a/gdeep/visualisation/visualiser.py
+++ b/gdeep/visualisation/visualiser.py
@@ -1,22 +1,25 @@
+from typing import List, Optional
+
 import torch
-from ..models import ModelExtractor
 from torchvision.utils import make_grid
-from gtda.plotting import plot_diagram, plot_betti_surfaces
-from . import persistence_diagrams_of_activations, \
-    plotly2tensor, Compactification, png2tensor
 from sklearn.manifold import MDS
 from gtda.diagrams import BettiCurve
 from html2image import Html2Image
 from captum.attr import visualization
-from matplotlib.colors import LinearSegmentedColormap
 import numpy as np
 import matplotlib.pyplot as plt
 from captum.attr import LayerAttribution
+from gtda.plotting import plot_diagram, plot_betti_surfaces
+from . import persistence_diagrams_of_activations, \
+    plotly2tensor, Compactification, png2tensor
+
+from gdeep.trainer import Trainer
+from gdeep.analysis.interpretability import Interpreter
+from gdeep.utility import DEVICE
+from ..models import ModelExtractor
 
-if torch.cuda.is_available():
-    DEVICE = torch.device("cuda")
-else:
-    DEVICE = torch.device("cpu")
+FONT_SIZE = 16
+Tensor = torch.Tensor
 
 
 class Visualiser:
@@ -25,13 +28,12 @@ class Visualiser:
     different kinds to tensors
 
     Args:
-        pipe (Pipeline):
+        pipe :
             the pipeline to get info from
-        writer (tensorboard SummaryWriter):
-            the tensorboard writer
+
     """
 
-    def __init__(self, pipe):
+    def __init__(self, pipe: Trainer):
         self.pipe = pipe
         self.persistence_diagrams = None
 
@@ -42,7 +44,7 @@ def plot_data_model(self):
         """
 
         dataiter = iter(self.pipe.dataloaders[0])
-        images, labels = dataiter.next()
+        images, labels = next(dataiter)
         # print(str(labels.item()))
         self.pipe.writer.add_graph(self.pipe.model, images.to(DEVICE))
         features_list = []
@@ -132,8 +134,8 @@ def plot_persistence_diagrams(self, example=None):
         list_of_dgms = []
         for i, persistence_diagram in enumerate(self.persistence_diagrams):
             plot_persistence_diagram = plot_diagram(persistence_diagram)
-            imgT = plotly2tensor(plot_persistence_diagram)
-            list_of_dgms.append(imgT)
+            img_t = plotly2tensor(plot_persistence_diagram)
+            list_of_dgms.append(img_t)
         features = torch.stack(list_of_dgms)
         # grid = make_grid(features)
         self.pipe.writer.add_images("persistence_diagrams_of_activations",
@@ -151,14 +153,14 @@ def plot_decision_boundary(self, compact=False):
         """
 
         me = ModelExtractor(self.pipe.model, self.pipe.loss_fn)
-        X = next(iter(self.pipe.dataloaders[0]))[0][0]
+        x = next(iter(self.pipe.dataloaders[0]))[0][0]
 
         if compact:
             # initlaisation of the compactification
             cc = Compactification(precision=0.1,
                                   n_samples=500,
                                   epsilon=0.051,
-                                  n_features=X.shape[0],
+                                  n_features=x.shape[0],
                                   n_epochs=100,
                                   neural_net=self.pipe.model)
 
@@ -172,24 +174,30 @@ def plot_decision_boundary(self, compact=False):
             self.pipe.writer.flush()
             return db, d_final, label_final
         else:
-            db = me.get_decision_boundary(X)
+            db = me.get_decision_boundary(x)
             self.pipe.writer.add_embedding(db,
                                            tag="decision_boundary",
                                            global_step=0)
             self.pipe.writer.flush()
             return db.cpu(), None, None
 
-    def betti_plot_layers(self, homology_dimension=(0, 1), example=None):
+    def betti_plot_layers(self, homology_dimension: List[int] = (0, 1),
+                          example: Optional[Tensor] = None):
         """
         Args:
-            homology_dimension (int array, optional):
+            homology_dimension :
                 An array of
                 homology dimensions
+            example :
+                the selected batch of data to
+                compute the activations on. By
+                defaults, this method takes the
+                first 100 elements
 
         Returns:
            (None):
                 a tuple of figures representing the Betti surfaces
-                of the data accross layers of the NN, with one figure
+                of the data across layers of the NN, with one figure
                 per dimension in homology_dimensions. Otherwise, a single
                 figure representing the Betti curve of the single sample
                 present.
@@ -202,22 +210,22 @@ def betti_plot_layers(self, homology_dimension=(0, 1), example=None):
                 inputs = inputs.to(example.dtype)
             self.persistence_diagrams = \
                 persistence_diagrams_of_activations(me.get_activations(inputs))
-        BC = BettiCurve()
-        dgms = BC.fit_transform(self.persistence_diagrams)
+        bc = BettiCurve()
+        dgms = bc.fit_transform(self.persistence_diagrams)
 
         plots = plot_betti_surfaces(dgms,
-                                    samplings=BC.samplings_,
+                                    samplings=bc.samplings_,
                                     homology_dimensions=homology_dimension)
 
         for i in range(len(plots)):
             plots[i].show()
 
-    def plot_interpreter_text(self, interpreter):
+    def plot_interpreter_text(self, interpreter: Interpreter):
         """This method allows to plot the results of an
         Interpreter for text data.
 
         Args:
-            interpreter (Interpreter):
+            interpreter :
                 this is a ``gdeep.analysis.interpretability``
                 initilised ``Interpreter`` class
                 
@@ -238,7 +246,7 @@ def plot_interpreter_text(self, interpreter):
                                    dataformats="HWC")
         return fig
 
-    def plot_interpreter_image(self, interpreter):
+    def plot_interpreter_image(self, interpreter: Interpreter):
         """This method allows to plot the results of an
         Interpreter for image data.
 
@@ -250,21 +258,18 @@ def plot_interpreter_image(self, interpreter):
         Returns:
             matplotlib.figure
         """
-        default_cmap = LinearSegmentedColormap.from_list('custom blue',
-                                                         [(0, '#ffffff'),
-                                                          (0.25, '#000000'),
-                                                          (1, '#000000')],
-                                                         N=256)
+
         try:
-            attrib = np.transpose(interpreter.attrib.squeeze().detach().cpu().numpy(),
-                                  (1, 2, 0))
+            attrib = torch.permute(interpreter.attribution.squeeze().detach(),
+                                   (1, 2, 0)).detach().cpu().numpy()
         except ValueError:
-            attrib = np.transpose(LayerAttribution.interpolate(interpreter.attrib.detach().cpu(),
-                                                               tuple(interpreter.image.squeeze().detach().cpu().shape[-2:])).squeeze(0),
-                                  (1, 2, 0))
-            attrib = torch.stack((attrib, attrib, attrib),axis=2).squeeze(-1).detach().cpu().numpy()
-        img = np.transpose(interpreter.image.squeeze().detach().cpu().numpy(),
-                           (1, 2, 0))
+            attrib = torch.permute(LayerAttribution.interpolate(interpreter.attribution.detach().cpu(),
+                                                                tuple(interpreter.image.squeeze().
+                                                                      detach().cpu().shape[-2:])).squeeze(0),
+                                   (1, 2, 0))
+            attrib = torch.stack([attrib, attrib, attrib], dim=2).squeeze(-1).detach().cpu().numpy()
+        img = torch.permute(interpreter.image.squeeze().detach(),
+                            (1, 2, 0)).detach().cpu().numpy()
 
         fig, _ = visualization.visualize_image_attr_multiple(
             attrib,
@@ -276,7 +281,7 @@ def plot_interpreter_image(self, interpreter):
                                     fig)
         return fig
 
-    def plot_interpreter_tabular(self, interpreter):
+    def plot_interpreter_tabular(self, interpreter: Interpreter):
         """This method allows to plot the results of an
         Interpreter for tabular data.
 
@@ -289,36 +294,12 @@ def plot_interpreter_tabular(self, interpreter):
             matplotlib.figure
         """
         # prepare attributions for visualization
-        X_test = interpreter.X
-        x_axis_data = np.arange(X_test.shape[1])
+        x_test = interpreter.x
+        x_axis_data = np.arange(x_test.shape[1])
         x_axis_data_labels = list(map(lambda idx: idx,
                                       x_axis_data))
-
-        ig_attr_test_sum = interpreter.ig_attr_test.detach().cpu().numpy().sum(0)
-        ig_attr_test_norm_sum = ig_attr_test_sum / \
-            np.linalg.norm(ig_attr_test_sum, ord=1)
-
-        ig_nt_attr_test_sum = \
-            interpreter.ig_nt_attr_test.detach().cpu().numpy().sum(0)
-        ig_nt_attr_test_norm_sum = ig_nt_attr_test_sum / \
-            np.linalg.norm(ig_nt_attr_test_sum, ord=1)
-
-        dl_attr_test_sum = \
-            interpreter.dl_attr_test.detach().cpu().numpy().sum(0)
-        dl_attr_test_norm_sum = dl_attr_test_sum / \
-            np.linalg.norm(dl_attr_test_sum, ord=1)
-
-        # gs_attr_test_sum = \
-        #     interpreter.gs_attr_test.detach().cpu().numpy().sum(0)
-        # gs_attr_test_norm_sum = gs_attr_test_sum / \
-        #     np.linalg.norm(gs_attr_test_sum, ord=1)
-
-        fa_attr_test_sum = interpreter.fa_attr_test.detach().cpu().numpy().sum(0)
-        fa_attr_test_norm_sum = fa_attr_test_sum / \
-            np.linalg.norm(fa_attr_test_sum, ord=1)
-
-        # lin_weight = model.lin1.weight[0].detach().numpy()
-        # y_axis_lin_weight = lin_weight / np.linalg.norm(lin_weight, ord=1)
+        attribution_sum = interpreter.attribution.detach().cpu().numpy().sum(0)
+        attribution_norm_sum = attribution_sum / np.linalg.norm(attribution_sum, ord=1)
 
         width = 0.14
         legends = ['Int Grads', 'Int Grads w/SmoothGrad', 'DeepLift',
@@ -327,28 +308,18 @@ def plot_interpreter_tabular(self, interpreter):
         fig = plt.figure(figsize=(20, 10))
 
         ax = plt.subplot()
-        ax.set_title('Comparing input feature importances across ' + \
-            'multiple algorithms and learned weights')
+        ax.set_title('Comparing input feature importances across ' +
+                     'multiple algorithms and learned weights')
         ax.set_ylabel('Attributions')
 
-        FONT_SIZE = 16
         plt.rc('font', size=FONT_SIZE)  # fontsize of the text sizes
         plt.rc('axes', titlesize=FONT_SIZE)  # fontsize of the axes title
         plt.rc('axes', labelsize=FONT_SIZE)  # fontsize of the x and y labels
         plt.rc('legend', fontsize=FONT_SIZE - 4)  # fontsize of the legend
 
-        ax.bar(x_axis_data, ig_attr_test_norm_sum, width, align='center',
+        ax.bar(x_axis_data, attribution_norm_sum, width, align='center',
                alpha=0.8, color='#eb5e7c')
-        ax.bar(x_axis_data + width, ig_nt_attr_test_norm_sum, width,
-               align='center', alpha=0.7, color='#A90000')
-        ax.bar(x_axis_data + 2 * width, dl_attr_test_norm_sum, width,
-               align='center', alpha=0.6, color='#34b8e0')
-        # ax.bar(x_axis_data + 3 * width, gs_attr_test_norm_sum,
-        # width, align='center',  alpha=0.8, color='#4260f5')
-        ax.bar(x_axis_data + 4 * width, fa_attr_test_norm_sum, width,
-               align='center', alpha=1.0, color='#49ba81')
-        # ax.bar(x_axis_data + 5 * width, y_axis_lin_weight, width,
-        # align='center', alpha=1.0, color='grey')
+
         ax.autoscale_view()
         plt.tight_layout()
 
diff --git a/requirements.txt b/requirements.txt
index 4ed34965..e1649098 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -24,6 +24,7 @@ html2image
 optuna
 sphinx_rtd_theme
 giotto-ph
+tqdm
 Pillow
 h5py
 dotmap
@@ -32,3 +33,4 @@ google-cloud-storage
 wget
 ipython
 jsonpickle
+typing_extensions; python_version == '3.7'


Test patch:
diff --git a/gdeep/analysis/decision_boundary/tests/test_compute_boundary.py b/gdeep/analysis/decision_boundary/tests/test_compute_topology.py
similarity index 100%
rename from gdeep/analysis/decision_boundary/tests/test_compute_boundary.py
rename to gdeep/analysis/decision_boundary/tests/test_compute_topology.py
diff --git a/gdeep/analysis/interpretability/tests/test_captum.py b/gdeep/analysis/interpretability/tests/test_captum.py
new file mode 100644
index 00000000..65a22bb2
--- /dev/null
+++ b/gdeep/analysis/interpretability/tests/test_captum.py
@@ -0,0 +1,77 @@
+from torch.utils.data.sampler import SubsetRandomSampler
+from torch import nn
+import torch
+
+from gdeep.data.datasets import DatasetBuilder
+from gdeep.analysis.interpretability import Interpreter
+from gdeep.visualisation import Visualiser
+from gdeep.data import PreprocessingPipeline
+from gdeep.data import TransformingDataset
+from gdeep.data.preprocessors import Normalization, TokenizerTextClassification
+from gdeep.data.datasets import DataLoaderBuilder
+from gdeep.trainer import Trainer
+from torch.utils.tensorboard import SummaryWriter
+
+
+writer = SummaryWriter()
+
+# many time we get an IterableDataset which is good for memory consumption, but cannot be sampled!
+# we can load the whole dataset in memory and sample it
+bd = DatasetBuilder(name="AG_NEWS",
+                    convert_to_map_dataset=True)
+ds_tr_str, ds_val_str, ds_ts_str = bd.build()
+ptd = TokenizerTextClassification()
+
+ptd.fit_to_dataset(ds_tr_str)
+transformed_textds = ptd.attach_transform_to_dataset(ds_tr_str)
+transformed_textts = ptd.attach_transform_to_dataset(ds_val_str)
+
+# the only part of the training/test set we are interested in
+train_indices = list(range(64*10))
+test_indices = list(range(64*5))
+
+dl_tr2, dl_ts2, _ = DataLoaderBuilder((transformed_textds,
+                                      transformed_textts)).build(({"batch_size":16,
+                                                                   "sampler":SubsetRandomSampler(train_indices)},
+                                                                  {"batch_size":16,
+                                                                   "sampler":SubsetRandomSampler(test_indices)}
+                                                                 ))
+
+class TextClassificationModel(nn.Module):
+
+    def __init__(self, vocab_size, embed_dim, num_class):
+        super(TextClassificationModel, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim, sparse=True)
+        self.fc = nn.Linear(embed_dim, num_class)
+        self.init_weights()
+
+    def init_weights(self):
+        initrange = 0.5
+        self.embedding.weight.data.uniform_(-initrange, initrange)
+        self.fc.weight.data.uniform_(-initrange, initrange)
+        self.fc.bias.data.zero_()
+
+    def forward(self, text):
+        embedded = self.embedding(text)
+        mean = torch.mean(embedded,dim=1)
+        return self.fc(mean)
+
+vocab_size = len(ptd.vocabulary)
+emsize = 64
+# print(vocab_size, emsize)
+model = TextClassificationModel(vocab_size, emsize, 4)
+loss_fn = nn.CrossEntropyLoss()
+pipe = Trainer(model, (dl_tr2, dl_ts2), loss_fn, writer)
+
+def test_interpret_text() -> None:
+    vs = Visualiser(pipe)
+    inter = Interpreter(pipe.model,
+                        method="LayerIntegratedGradients")
+
+    inter.interpret_text("I am writing about money and business",
+                         0,
+                         ptd.vocabulary,
+                         ptd.tokenizer,
+                         layer=pipe.model.embedding)
+
+    vs.plot_interpreter_text(inter)
\ No newline at end of file
diff --git a/gdeep/models/tests/test_extractor.py b/gdeep/models/tests/test_extractor.py
new file mode 100644
index 00000000..dad84ffc
--- /dev/null
+++ b/gdeep/models/tests/test_extractor.py
@@ -0,0 +1,55 @@
+"""Testing for simpleNN."""
+# License: GNU AGPLv3
+
+import torch
+from torch import nn
+
+from gdeep.models import ModelExtractor
+from ..simple_nn import FFNet
+from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder
+
+
+bd = DatasetBuilder(name="Blobs")
+ds_tr, ds_val, _ = bd.build()
+dl = DataLoaderBuilder([ds_tr, ds_val])
+dl_tr, dl_val, dl_ts = dl.build()  # sampler=SubsetRandomSampler(train_indices))
+
+
+class Model1(nn.Module):
+    def __init__(self):
+        super(Model1, self).__init__()
+        self.seqmodel = nn.Sequential(nn.Flatten(), FFNet(arch=[3, 5, 10, 5, 2]))
+
+    def forward(self, x):
+        return self.seqmodel(x)
+
+
+def test_extractor_get_layers_param():
+    model = Model1()
+
+    loss_fn = nn.CrossEntropyLoss()
+    me = ModelExtractor(model, loss_fn)
+
+    lista = me.get_layers_param()
+
+    for k, item in lista.items():
+        k, item.shape  # noqa
+
+
+def test_extractor_get_decision_boundary():
+    model = FFNet(arch=[3, 3])
+    loss_fn = nn.CrossEntropyLoss()
+    me = ModelExtractor(model, loss_fn)
+    x = next(iter(dl_tr))[0][0]
+    if x.dtype is not torch.int64:
+        res = me.get_decision_boundary(x, n_epochs=5)
+        assert res.shape[-1] == 3
+
+
+def test_extractor_get_activations():
+    model = FFNet(arch=[3, 3])
+    loss_fn = nn.CrossEntropyLoss()
+    me = ModelExtractor(model, loss_fn)
+    x = next(iter(dl_tr))[0][0]
+    list_activations = me.get_activations(x)
+    assert len(list_activations) == 2
diff --git a/gdeep/models/tests/test_periodic_nn.py b/gdeep/models/tests/test_periodic_nn.py
new file mode 100644
index 00000000..d25565b5
--- /dev/null
+++ b/gdeep/models/tests/test_periodic_nn.py
@@ -0,0 +1,14 @@
+"""Testing for simpleNN."""
+# License: GNU AGPLv3
+import pytest
+import torch
+from .. import FFNet, PeriodicNeuralNetwork
+
+
+def test_periodic():
+    X = torch.rand(10, 2)
+    sn = FFNet([2, 3, 2])
+    bdry_list = [(-1,1),(0.5,1)]
+    period = PeriodicNeuralNetwork(sn, bdry_list)
+    out = period.forward(X)
+    assert out.shape[0] == 10 and out.shape[1] == 2
diff --git a/gdeep/models/tests/test_simple_nn.py b/gdeep/models/tests/test_simple_nn.py
index 9a4392ff..de4929ce 100644
--- a/gdeep/models/tests/test_simple_nn.py
+++ b/gdeep/models/tests/test_simple_nn.py
@@ -9,5 +9,4 @@ def test_FFNet():
     X = torch.rand(10, 2)
     sn = FFNet([2, 3, 2])
     out = sn.forward(X)
-    print(out.shape)
     assert out.shape[0] == 10 and out.shape[1] == 2
diff --git a/gdeep/search/tests/test_benchmark.py b/gdeep/search/tests/test_benchmark.py
new file mode 100644
index 00000000..f3fcc048
--- /dev/null
+++ b/gdeep/search/tests/test_benchmark.py
@@ -0,0 +1,113 @@
+from torch import nn
+import torchvision.models as models
+from torch.optim import SGD
+from torch.utils.data.sampler import SubsetRandomSampler
+
+from gdeep.models import FFNet
+from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder
+from gdeep.data.preprocessors import ToTensorImage
+from gdeep.search import HyperParameterOptimization, GiottoSummaryWriter, Benchmark
+
+
+# parametric model with string value
+class model2(nn.Module):
+    def __init__(self, n_nodes="100"):
+        super(model2, self).__init__()
+        self.md = nn.Sequential(nn.Sequential(models.resnet18(pretrained=True),
+                                              nn.Linear(1000, eval(n_nodes))),
+                                nn.Linear(eval(n_nodes), 10))
+
+    def forward(self, x):
+        return self.md(x)
+
+
+models_dicts = []
+
+model = model2()
+
+temp_dict = {}
+temp_dict["name"] = "resnet18"
+temp_dict["model"] = model
+
+models_dicts.append(temp_dict)
+
+
+# avoid having exposed parameters that wll not be searched on
+class model_no_param(nn.Module):
+    def __init__(self):
+        super(model_no_param, self).__init__()
+        self.mod = FFNet([3, 5, 5, 2])
+
+    def forward(self, x):
+        return self.mod(x)
+
+
+model5 = model_no_param()
+temp_dict = {}
+temp_dict["name"] = "ffnn"
+temp_dict["model"] = model5
+
+models_dicts.append(temp_dict)
+
+
+writer = GiottoSummaryWriter()
+
+# download CIFAR10
+bd = DatasetBuilder(name="CIFAR10")
+ds_tr, _, _ = bd.build(download=True)
+
+# Preprocessing
+transformation = ToTensorImage((32, 32))
+transformation.fit_to_dataset(ds_tr)  # this is useless for this transformation
+
+transformed_ds_tr = transformation.attach_transform_to_dataset(ds_tr)
+
+dataloaders_dicts = []
+
+
+test_indices = [64*5 + x for x in range(32*3)]
+train_indices = [x for x in range(32*2)]
+
+dl = DataLoaderBuilder((transformed_ds_tr, transformed_ds_tr))
+dl_tr, dl_val, _ = dl.build(({"batch_size":32,"sampler":SubsetRandomSampler(train_indices)},
+                             {"batch_size":32,"sampler":SubsetRandomSampler(test_indices)}))
+
+
+temp_dict = {}
+temp_dict["name"] = "CIFAR10_1000"
+temp_dict["dataloaders"] = (dl_tr, dl_val, _)
+
+dataloaders_dicts.append(temp_dict)
+
+db = DatasetBuilder(name="DoubleTori")
+ds_tr, ds_val, _ = db.build()
+
+dl_tr, dl_ts, _ = DataLoaderBuilder((ds_tr, ds_val)).build(({"batch_size":48},{"batch_size":32}))
+
+temp_dict = {}
+temp_dict["name"] = "double_tori"
+temp_dict["dataloaders"] = (dl_tr, dl_ts)
+
+dataloaders_dicts.append(temp_dict)
+
+
+def test_benchmark():
+    # initialise loss
+    loss_fn = nn.CrossEntropyLoss()
+
+    # initialise the benchmarking class. When we do not specify it, it will use KFold with 5 splits
+    bench = Benchmark(models_dicts, dataloaders_dicts, loss_fn, writer)
+
+    # start the benchmarking
+    bench.start(SGD, 1, True, {"lr": 0.01}, {"batch_size": 23})
+
+
+def test_benchmark2():
+    # initialise loss
+    loss_fn = nn.CrossEntropyLoss()
+
+    # initialise the benchmarking class. When we do not specify it, it will use KFold with 5 splits
+    bench = Benchmark(models_dicts, dataloaders_dicts, loss_fn, writer)
+
+    # start the benchmarking
+    bench.start(SGD, 2, False, {"lr": 0.01}, {"batch_size": 32}, n_accumulated_grads=2)
diff --git a/gdeep/search/tests/test_hpo.py b/gdeep/search/tests/test_hpo.py
new file mode 100644
index 00000000..364efbd5
--- /dev/null
+++ b/gdeep/search/tests/test_hpo.py
@@ -0,0 +1,157 @@
+from torch import nn
+import torchvision.models as models
+from torch.optim import SGD, Adam
+from sklearn.model_selection import StratifiedKFold
+from torch.utils.data.sampler import SubsetRandomSampler
+
+from gdeep.trainer import Trainer
+from gdeep.data.datasets import DatasetBuilder, DataLoaderBuilder
+from gdeep.data.preprocessors import ToTensorImage
+from gdeep.search import HyperParameterOptimization, GiottoSummaryWriter
+
+
+# parametric model with string value
+class Model2(nn.Module):
+    def __init__(self, n_nodes="100"):
+        super(Model2, self).__init__()
+        self.md = nn.Sequential(nn.Sequential(models.resnet18(pretrained=True),
+                                              nn.Linear(1000, eval(n_nodes))),
+                                nn.Linear(eval(n_nodes), 10))
+
+    def forward(self, x):
+        return self.md(x)
+
+
+writer = GiottoSummaryWriter()
+
+# download MNIST
+bd = DatasetBuilder(name="CIFAR10")
+ds_tr, _, _ = bd.build(download=True)
+
+# Preprocessing
+transformation = ToTensorImage([32, 32])
+transformation.fit_to_dataset(ds_tr)  # this is useless for this transformation
+
+transformed_ds_tr = transformation.attach_transform_to_dataset(ds_tr)
+
+# use only 32*7 images from CIFAR10
+train_indices = list(range(32 * 7))
+dl_tr, *_ = DataLoaderBuilder([transformed_ds_tr]).build(
+    [{"batch_size": 32, "sampler": SubsetRandomSampler(train_indices)}])
+
+
+def test_hpo_failure():
+    # define the model
+    model = Model2()
+
+    # initialise loss
+    loss_fn = nn.CrossEntropyLoss()
+
+    # initialise pipeline class
+    pipe = Trainer(model, [dl_tr, None], loss_fn, writer,
+                   k_fold_class=StratifiedKFold(2, shuffle=True))
+
+    # initialise gridsearch
+    try:
+        HyperParameterOptimization(pipe, "accu", 2, best_not_last=True)
+    except AssertionError:
+        pass
+
+def test_hpo_cross_val():
+    # define the model
+    model = Model2()
+
+    # initialise loss
+    loss_fn = nn.CrossEntropyLoss()
+
+    # initialise pipeline class
+    pipe = Trainer(model, [dl_tr, None], loss_fn, writer,
+                   k_fold_class=StratifiedKFold(2, shuffle=True))
+
+    # initialise gridsearch
+    search = HyperParameterOptimization(pipe, "accuracy", 2, best_not_last=True)
+
+    # if you want to store pickle files of the models instead of the state_dicts
+    search.store_pickle = True
+
+    # dictionaries of hyperparameters
+    optimizers_params = {"lr": [0.001, 0.01]}
+    dataloaders_params = {"batch_size": [32, 64, 16]}
+    models_hyperparams = {"n_nodes": ["200"]}
+
+    # starting the gridsearch
+    search.start([SGD, Adam], 3, False, optimizers_params, dataloaders_params,
+                 models_hyperparams, n_accumulated_grads=2)
+
+
+def test_hpo_accumulated_grads():
+    # define the model
+    model = Model2()
+
+    # initialise loss
+    loss_fn = nn.CrossEntropyLoss()
+
+    # initialise pipeline class
+    pipe = Trainer(model, [dl_tr, None], loss_fn, writer)
+
+    # initialise gridsearch
+    search = HyperParameterOptimization(pipe, "accuracy", 2, best_not_last=True)
+
+    # if you want to store pickle files of the models instead of the state_dicts
+    search.store_pickle = True
+
+    # dictionaries of hyperparameters
+    optimizers_params = {"lr": [0.001, 0.01]}
+    dataloaders_params = {"batch_size": [32, 64, 16]}
+    models_hyperparams = {"n_nodes": ["200"]}
+
+    # starting the gridsearch
+    search.start([SGD, Adam], 2, False, optimizers_params, dataloaders_params,
+                 models_hyperparams, n_accumulated_grads=2)
+
+
+def test_hpo_loss():
+    # define the model
+    model = Model2()
+
+    # initialise loss
+    loss_fn = nn.CrossEntropyLoss()
+
+    # initialise pipeline class
+    pipe = Trainer(model, [dl_tr, None], loss_fn, writer)
+
+    # initialise gridsearch
+    search = HyperParameterOptimization(pipe, "loss", 2, best_not_last=True)
+
+    # dictionaries of hyperparameters
+    optimizers_params = {"lr": [0.001, 0.01]}
+    dataloaders_params = {"batch_size": [32, 64, 16]}
+    models_hyperparams = {"n_nodes": ["200"]}
+
+    # starting the gridsearch
+    search.start([SGD, Adam], 2, False, optimizers_params, dataloaders_params, models_hyperparams)
+
+
+def test_hpo_string_parameters():
+    # define the model
+    model = Model2()
+
+    # initialise loss
+    loss_fn = nn.CrossEntropyLoss()
+
+    # initialise pipeline class
+    pipe = Trainer(model, [dl_tr, None], loss_fn, writer)
+
+    # initialise gridsearch
+    search = HyperParameterOptimization(pipe, "loss", 2)
+
+    # if you want to store pickle files of the models instead of the state_dicts
+    search.store_pickle = True
+
+    # dictionaries of hyperparameters
+    optimizers_params = {"lr": [0.001, 0.01]}
+    dataloaders_params = {"batch_size": [32, 64]}
+    models_hyperparams = {"n_nodes": ["200", "256"]}
+
+    # starting the gridsearch
+    search.start([SGD], 1, False, optimizers_params, dataloaders_params, models_hyperparams)
diff --git a/gdeep/search/tests/test_hpo_config.py b/gdeep/search/tests/test_hpo_config.py
new file mode 100644
index 00000000..7af369df
--- /dev/null
+++ b/gdeep/search/tests/test_hpo_config.py
@@ -0,0 +1,9 @@
+
+from gdeep.search import HPOConfig
+from torch.optim import SGD
+
+
+def test_hpo_config() -> None:
+    config = HPOConfig([SGD])
+    dictionary = config.to_dict()
+    assert isinstance(dictionary, dict)
diff --git a/gdeep/trainer/tests/test_metrics.py b/gdeep/trainer/tests/test_metrics.py
new file mode 100644
index 00000000..20f92a5e
--- /dev/null
+++ b/gdeep/trainer/tests/test_metrics.py
@@ -0,0 +1,12 @@
+import torch
+
+from gdeep.trainer import accuracy
+
+
+def test_accuracy():
+    predictions = torch.rand((10, 3))
+    y = torch.randint(0, 2, (10,))
+
+    acc = accuracy(predictions, y)
+
+    assert acc >= 0
diff --git a/gdeep/trainer/tests/test_trainer.py b/gdeep/trainer/tests/test_trainer.py
index eca84d62..05216eb9 100644
--- a/gdeep/trainer/tests/test_trainer.py
+++ b/gdeep/trainer/tests/test_trainer.py
@@ -1,6 +1,5 @@
 from torch import nn
 from torch.optim import SGD
-from sklearn.model_selection import StratifiedKFold
 import numpy as np
 
 from gdeep.trainer import Trainer
@@ -35,6 +34,6 @@ def forward(self, x):
     writer = GiottoSummaryWriter()
     # pipeline
     pipe = Trainer(model, [dl_tr, None],
-                    loss_fn, writer)#,StratifiedKFold(5, shuffle=True))
+                   loss_fn, writer)
     # then one needs to train the model using the pipeline!
     pipe.train(SGD, 2, True, {"lr": 0.001}, n_accumulated_grads=2)
diff --git a/gdeep/trainer/tests/test_trainer_config.py b/gdeep/trainer/tests/test_trainer_config.py
new file mode 100644
index 00000000..7c7aa5d8
--- /dev/null
+++ b/gdeep/trainer/tests/test_trainer_config.py
@@ -0,0 +1,9 @@
+
+from gdeep.trainer import TrainerConfig
+from torch.optim import SGD
+
+
+def test_trainer_config() -> None:
+    config = TrainerConfig(SGD)
+    dictionary = config.to_dict()
+    assert isinstance(dictionary, dict)
diff --git a/gdeep/utility/optimisation/tests/test_persistence_grad.py b/gdeep/utility/optimisation/tests/test_persistence_grad.py
index 9e798dc0..9b163516 100644
--- a/gdeep/utility/optimisation/tests/test_persistence_grad.py
+++ b/gdeep/utility/optimisation/tests/test_persistence_grad.py
@@ -17,7 +17,7 @@ def test_PersistenceGradient_2d():
     assert pg.phi(X).shape[0] == 14
 
     assert pg.persistence_function(X).item() >= -2.328427314758301 - 0.001
-    pg.SGD(X, n_epochs=4, lr=0.4)
+    pg.sgd(X, n_epochs=4, lr=0.4)
 
 
 def test_PersistenceGradient_3d():
@@ -33,7 +33,7 @@ def test_PersistenceGradient_3d():
     assert pg.phi(X).shape[0] == 14
 
     assert pg.persistence_function(X).item() >= -2.7783148288726807 - 0.001
-    pg.SGD(X, n_epochs=4, lr=0.4)
+    pg.sgd(X, n_epochs=4, lr=0.4)
 
 
 def test_PersistenceGradient_5d():
@@ -49,7 +49,7 @@ def test_PersistenceGradient_5d():
     assert X.dtype == torch.float32
     assert pg.phi(X).shape[0] == 14
     assert pg.persistence_function(X).item() >= -2.7783148288726807 - 0.001
-    pg.SGD(X, n_epochs=4, lr=0.4)
+    pg.sgd(X, n_epochs=4, lr=0.4)
 
 
 def test_PersistenceGradient_4d():
@@ -65,7 +65,7 @@ def test_PersistenceGradient_4d():
     assert X.dtype == torch.float32
     assert pg.phi(X).shape[0] == 14
     assert pg.persistence_function(X).item() >= 0.3467579483985901 + 0.001
-    pg.SGD(X)
+    pg.sgd(X)
 
 
 def test_PersistenceGradient_matrix():
@@ -80,7 +80,7 @@ def test_PersistenceGradient_matrix():
                              metric="precomputed")
     assert all(pg.phi(dist) == torch.tensor([0., 0., 0., 2., 2.2, 3., 3.]))
     assert pg.persistence_function(dist).item() >= -6.3 - 0.0001
-    fig, fig3d, loss_val = pg.SGD(dist, n_epochs=1,
+    fig, fig3d, loss_val = pg.sgd(dist, n_epochs=1,
                                   lr=1)
 
 
@@ -98,7 +98,7 @@ def test_PersistenceGradient_matrix_2():
                                              1., 2., 2., 2., 10.,
                                              10., 10., 10., 10.]))
     assert pg.persistence_function(dist).item() > -23.
-    fig, fig3d, loss_val = pg.SGD(dist, n_epochs=1,
+    fig, fig3d, loss_val = pg.sgd(dist, n_epochs=1,
                                   lr=0.002)
                                   
 
@@ -115,7 +115,7 @@ def test_PersistenceGradient_pts():
     assert all(pg.phi(pts)[:5] == torch.tensor([0., 0., 0.,
                                                 1., 1.]))
 
-    fig, fig3d, loss_val = pg.SGD(pts, n_epochs=1,
+    fig, fig3d, loss_val = pg.sgd(pts, n_epochs=1,
                                   lr=0.002)
 
     assert (pts.grad == torch.tensor([[ 1.,  1.],
diff --git a/gdeep/visualisation/tests/test_compactification.py b/gdeep/visualisation/tests/test_compactification.py
new file mode 100644
index 00000000..a82feb5f
--- /dev/null
+++ b/gdeep/visualisation/tests/test_compactification.py
@@ -0,0 +1,16 @@
+
+from gdeep.visualisation import Compactification
+from gdeep.models import FFNet
+
+
+def test_compactification():
+    model = FFNet([2, 4, 2])
+    c = Compactification(0.4,
+                         2,
+                         1000,
+                         0.05,
+                         50,
+                         [(-1., 1.), (-1., 1.)],
+                         model)
+    c.create_final_distance_matrix()
+    c.plot_chart(1)
diff --git a/gdeep/visualisation/tests/test_visualiser.py b/gdeep/visualisation/tests/test_visualiser.py
new file mode 100644
index 00000000..4074c78b
--- /dev/null
+++ b/gdeep/visualisation/tests/test_visualiser.py
@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+from torch.utils.data.sampler import SubsetRandomSampler
+from torch.utils.tensorboard import SummaryWriter
+
+from gdeep.visualisation import Visualiser
+from gdeep.data.preprocessors import TokenizerTextClassification
+from gdeep.data.datasets import DataLoaderBuilder, DatasetBuilder
+from gdeep.trainer import Trainer
+
+bd = DatasetBuilder(name="AG_NEWS", convert_to_map_dataset=True)
+ds_tr_str, ds_val_str, ds_ts_str = bd.build()
+
+
+ptd = TokenizerTextClassification()
+
+ptd.fit_to_dataset(ds_tr_str)
+transformed_textds = ptd.attach_transform_to_dataset(ds_tr_str)
+transformed_textts = ptd.attach_transform_to_dataset(ds_val_str)
+
+# the only part of the training/test set we are interested in
+train_indices = list(range(64*10))
+test_indices = list(range(64*5))
+
+dl_tr2, dl_ts2, _ = DataLoaderBuilder([transformed_textds,
+                                      transformed_textts]).build([{"batch_size": 16,
+                                                                   "sampler": SubsetRandomSampler(train_indices)},
+                                                                  {"batch_size": 16,
+                                                                   "sampler": SubsetRandomSampler(test_indices)}
+                                                                  ])
+writer = SummaryWriter()
+
+
+class TextClassificationModel(nn.Module):
+
+    def __init__(self, vocab_size, embed_dim, num_class):
+        super(TextClassificationModel, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim, sparse=True)
+        self.fc = nn.Linear(embed_dim, num_class)
+        self.init_weights()
+
+    def init_weights(self):
+        initrange = 0.5
+        self.embedding.weight.data.uniform_(-initrange, initrange)
+        self.fc.weight.data.uniform_(-initrange, initrange)
+        self.fc.bias.data.zero_()
+
+    def forward(self, text):
+        embedded = self.embedding(text)
+        mean = torch.mean(embedded, dim=1)
+        return self.fc(mean)
+
+
+def test_visualiser():
+    vocab_size = len(ptd.vocabulary)
+    emsize = 64
+    loss_fn = nn.CrossEntropyLoss()
+    model = TextClassificationModel(vocab_size, emsize, 4)
+    pipe = Trainer(model, [dl_tr2, dl_ts2], loss_fn, writer)
+
+    vs = Visualiser(pipe)
+
+    x, _ = next(iter(dl_tr2))
+    vs.plot_data_model()
+    vs.plot_activations(x)
+    vs.plot_persistence_diagrams(x)
+    vs.betti_plot_layers([0, 1], x)