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

Issue:
Have consistent code formatting.
Description: Different parts of the code will use different conventions for
variable and function names, comments, and general organization. This
leads to inconsistency which can make code difficult to read.

Resolution: Follow a style guide so that all code has a consistent appearance.
Agree on a style guide with all developers in our team and define it in a config file.

Tools: Use Black (https://github.com/psf/black) for code formatting. It can be
set up to run automatically on commit.
[Proposal] Add `pre-commit`
Hey maintainers 👋🏽 ☕ , I noticed that most of (if not all) the steps of the **[Python Coding Style](https://giotto-ai.github.io/gtda-docs/latest/contributing/#python-coding-style)** can be automated via [pre-commit](https://pre-commit.com/) such as `flake8` and `pytest` or even the ones in the GH actions such as `mypy`. Having [pre-commit](https://pre-commit.com/) would make that process much easier and enforce that the developers follow coding conventions. Maybe this could even help with #67 if we add `black` as a hook

**Additional context**
Although this would add `pre-commit` as a dependency for the project.



Discussion / hints:
also performed the refactoring with `jblack` https://github.com/irahorecka/jupyterblack
We shall add the use of black as code formatter in the issue template, but I would not force the refactoring on the merge.


Fix patch:
diff --git a/.github/workflows/deploy-gh-pages.yml b/.github/workflows/deploy-gh-pages.yml
index a8725cac..47e6988d 100644
--- a/.github/workflows/deploy-gh-pages.yml
+++ b/.github/workflows/deploy-gh-pages.yml
@@ -39,4 +39,3 @@ jobs:
         git add -f docs/build
         git commit -m "push sphinx build"
         git push origin gh-pages
-
diff --git a/.gitignore b/.gitignore
index 060291d7..b8f2a459 100755
--- a/.gitignore
+++ b/.gitignore
@@ -93,4 +93,4 @@ state_dicts/*
 # pylance files
 *.pyi
 
-giotto-env/
\ No newline at end of file
+giotto-env/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 00000000..3d4376ac
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,16 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.3.0
+    hooks:
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+    -   id: check-yaml
+    -   id: check-added-large-files
+-   repo: https://github.com/psf/black
+    rev: 22.6.0
+    hooks:
+    -   id: black
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: 'v0.971'  # Use the sha / tag you want to point at
+    hooks:
+    -   id: mypy
diff --git a/GOVERNANCE.rst b/GOVERNANCE.rst
index a974b857..0a3e2842 100644
--- a/GOVERNANCE.rst
+++ b/GOVERNANCE.rst
@@ -18,4 +18,3 @@ Giotto-deep Project Team:
 
 Former Project Team Members:
 ----------------------------
-
diff --git a/LICENSE b/LICENSE
index 867227df..37e986f4 100644
--- a/LICENSE
+++ b/LICENSE
@@ -7,8 +7,8 @@ You may obtain a copy of the License below or at https://www.gnu.org/licenses/ag
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
-   limitations under the License.                    
-                    
+   limitations under the License.
+
                     GNU AFFERO GENERAL PUBLIC LICENSE
                        Version 3, 19 November 2007
 
diff --git a/PULL_REQUEST_TEMPLATE.md b/PULL_REQUEST_TEMPLATE.md
index 21041aad..53f7380c 100644
--- a/PULL_REQUEST_TEMPLATE.md
+++ b/PULL_REQUEST_TEMPLATE.md
@@ -34,7 +34,7 @@ Describe your changes in detail.
 
 **Checklist**
 <!--
-Go over all the following points, and put an `x` in all the boxes that apply. 
+Go over all the following points, and put an `x` in all the boxes that apply.
 If you're unsure about any of these, don't hesitate to ask. We're here to help!
 -->
 - [ ] I have read the [guidelines for contributing](https://giotto-ai.github.io/gtda-docs/latest/contributing/#guidelines).
diff --git a/README.md b/README.md
index 83713954..324070eb 100644
--- a/README.md
+++ b/README.md
@@ -34,7 +34,7 @@ The simplest way to install `giotto-deep` is using `pip`:
 ```
 python -m pip install -U giotto-deep
 ```
-If necessary, this command will also automatically install all the library dependencies. 
+If necessary, this command will also automatically install all the library dependencies.
 **Note:** we recommend upgrading ``pip`` to a recent version as the above may fail on very old versions.
 
 
@@ -45,7 +45,7 @@ The first step to install the developer version of the package is to `git clone`
 git clone https://github.com/giotto-ai/giotto-deep.git
 ```
 The change the current working directory to the Repository root folder, e.g. `cd giotto-deep`.
-Make sure you have the latest version of pytorch installed. 
+Make sure you have the latest version of pytorch installed.
 You can do this by running the following command (if you have a GPU):
 ```
 pip install torch --extra-index-url https://download.pytorch.org/whl/cu113
@@ -59,7 +59,17 @@ Make sure you have upgraded to the last version of `pip` with
 python -m pip install --upgrade pip
 ```
 
-## Run local tests
+
+## Contributing
+The `giotto-deep` project welcomes contributions of all kinds. Please see our [contributing guidelines](
+    https://giotto-ai.github.io/gtda-docs/latest/contributing/#guidelines
+) for more information.
+
+We are using pre-commit hooks to ensure that the code is formatted correctly. To install the pre-commit hooks, run the following command from the root folder:
+```
+pre-commit install
+```
+
 To run both unit and integration tests on *macOS* or *Linux*, simply run the following command from the root folder:
 ```
 bash local_test.bh
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 80861f38..9900ab0e 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -14,17 +14,17 @@
 import sys
 import sphinx_rtd_theme  # noqa
 
-sys.path.insert(0, os.path.abspath(os.path.join('..', '..')))
+sys.path.insert(0, os.path.abspath(os.path.join("..", "..")))
 
 
 # -- Project information -----------------------------------------------------
 
-project = 'giotto-deep'
-copyright = '2022, L2F SA'  # noqa
-author = 'Matteo Caorsi'
+project = "giotto-deep"
+copyright = "2022, L2F SA"  # noqa
+author = "Matteo Caorsi"
 
 # The full version, including alpha/beta/rc tags
-release = '0.0.1'
+release = "0.0.1"
 
 
 # -- General configuration ---------------------------------------------------
@@ -32,11 +32,10 @@
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
-extensions = ['sphinx.ext.autodoc',
-              'sphinx_rtd_theme']
+extensions = ["sphinx.ext.autodoc", "sphinx_rtd_theme"]
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ["_templates"]
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
@@ -49,7 +48,7 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = 'sphinx_rtd_theme'
+html_theme = "sphinx_rtd_theme"
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
diff --git a/docs/source/distributed.rst b/docs/source/distributed.rst
index 5fab99c9..1a16e4b3 100644
--- a/docs/source/distributed.rst
+++ b/docs/source/distributed.rst
@@ -21,5 +21,3 @@ Starting distributed computations
 *********************************
 
 The distribution of computations is done automatically, you just to wrap all your HPO details inside a single file. In particular, in the ``/kubernetes/examples/`` folder you will find a notebook and a ``.py`` file. The ``.py`` file contains all the HPO relevant code, which you will have to customise following your needs. The notebook simply launches the enqueueing (and hence the jobs will run on the workers pods).
-
-
diff --git a/docs/source/installation.rst b/docs/source/installation.rst
index 5bbd0060..dc247655 100644
--- a/docs/source/installation.rst
+++ b/docs/source/installation.rst
@@ -32,7 +32,7 @@ The simplest way to install ``giotto-deep`` is using ``pip``   ::
 
     python -m pip install -U giotto-deep
 
-If necessary, this command will also automatically install all the above dependencies. 
+If necessary, this command will also automatically install all the above dependencies.
 Note: we recommend upgrading ``pip`` to a recent version as the above may fail on very old versions.
 
 **********************
@@ -69,7 +69,7 @@ After installation, you can launch the test suite from inside the
 ``/gdeep`` directory::
 
     pytest .
-    
+
 You can also run the typing checks from ``/gdeep`` using::
 
    mypy --ignore-missing-imports .
@@ -77,6 +77,3 @@ You can also run the typing checks from ``/gdeep`` using::
 You can also run the bash script for local tests (from the root folder!) like this::
 
    bash local_test.bh
-
-
-
diff --git a/docs/source/introduction.rst b/docs/source/introduction.rst
index 25e96a34..d22474ff 100644
--- a/docs/source/introduction.rst
+++ b/docs/source/introduction.rst
@@ -22,7 +22,7 @@ sections will explain the most important ones.
 DatasetBuilder and DataLoaderBuilder
 ====================================
 
-You can easily create or load datasets using the ``DatasetBuilder`` class, apply preprocessing and 
+You can easily create or load datasets using the ``DatasetBuilder`` class, apply preprocessing and
 then get them ready to train using the ``DataLoaderBuilder`` class.
 
 .. code-block:: python
@@ -39,7 +39,7 @@ then get them ready to train using the ``DataLoaderBuilder`` class.
 Model
 ======
 
-Models are vanilla ``torch.nn.Module`` and subclasses of it. You have all the freedom of PyTorch to 
+Models are vanilla ``torch.nn.Module`` and subclasses of it. You have all the freedom of PyTorch to
 Build your custom model.
 
 Trainer
@@ -47,7 +47,7 @@ Trainer
 
 Once you have setup your data, your model, your loss function (and maybe a performance metric), you
 can put all these ingredients together into the  ``Trainer`` class. Then, with the ``train`` method
-you can run the training and validation of your models, specifying all the parameters the you need. 
+you can run the training and validation of your models, specifying all the parameters the you need.
 
 .. code-block:: python
 
@@ -60,8 +60,8 @@ you can run the training and validation of your models, specifying all the param
 HyperPatameterOptimisation
 ==========================
 
-Instead of a single training, you may also want to search the space of hyper parameters to find the 
-Best model possible. In ``giotto-deep`` this step can be done in a few lines using the class 
+Instead of a single training, you may also want to search the space of hyper parameters to find the
+Best model possible. In ``giotto-deep`` this step can be done in a few lines using the class
 ``HyperPatameterOptimisation``:
 
 .. code-block:: python
@@ -91,8 +91,8 @@ of the ``Interpreter`` class.
 .. code-block:: python
 
    inter = Interpreter(pipe.model, method="LayerIntegratedGradients")
-   inter.interpret_text("I am writing about money and business", 
-       0, 
+   inter.interpret_text("I am writing about money and business",
+       0,
        ptd.vocabulary,
        ptd.tokenizer,
        layer=pipe.model.embedding,
@@ -112,4 +112,3 @@ or the heat maps resulting form the `Interpreter``: you need the ``Visualiser``
    vs = Visualiser(pipe)
    vs.plot_interpreter_text(inter)
    vs.plot_data_model()
-
diff --git a/docs/source/modules/data.rst b/docs/source/modules/data.rst
index 7299e43c..6d8163d6 100644
--- a/docs/source/modules/data.rst
+++ b/docs/source/modules/data.rst
@@ -3,7 +3,7 @@ Data
 
 .. automodule:: gdeep.data
    :members:
-   
+
 Preprocessors
 -------------
 
diff --git a/docs/source/modules/index.rst b/docs/source/modules/index.rst
index 6fb1204f..6cc79027 100644
--- a/docs/source/modules/index.rst
+++ b/docs/source/modules/index.rst
@@ -7,13 +7,13 @@ This pages contains a list of available features in the library.
    :maxdepth: 3
 
    analysis
-   
+
    data
-   
+
    models
-   
+
    trainer
-   
+
    search
 
    topology_layers
diff --git a/docs/source/modules/models.rst b/docs/source/modules/models.rst
index 4290d3bd..9efe3158 100644
--- a/docs/source/modules/models.rst
+++ b/docs/source/modules/models.rst
@@ -3,4 +3,3 @@ Models
 
 .. automodule:: gdeep.models
    :members:
-   
diff --git a/docs/source/modules/search.rst b/docs/source/modules/search.rst
index 185260b8..41baacf2 100644
--- a/docs/source/modules/search.rst
+++ b/docs/source/modules/search.rst
@@ -3,4 +3,3 @@ Search
 
 .. automodule:: gdeep.search
    :members:
-   
diff --git a/docs/source/modules/topology_layers.rst b/docs/source/modules/topology_layers.rst
index 3a28121a..b962370f 100644
--- a/docs/source/modules/topology_layers.rst
+++ b/docs/source/modules/topology_layers.rst
@@ -14,4 +14,4 @@ Pooling Layers
 --------------
 
 .. automodule:: gdeep.gdeep.topology_layers.pooling_layers
-   :members:   
+   :members:
diff --git a/docs/source/modules/trainer.rst b/docs/source/modules/trainer.rst
index f4dee4ed..3dcb22ef 100644
--- a/docs/source/modules/trainer.rst
+++ b/docs/source/modules/trainer.rst
@@ -3,4 +3,3 @@ Trainer
 
 .. automodule:: gdeep.trainer
    :members:
-   
diff --git a/docs/source/modules/utility.rst b/docs/source/modules/utility.rst
index 8bf756a4..0c3d9df1 100644
--- a/docs/source/modules/utility.rst
+++ b/docs/source/modules/utility.rst
@@ -3,7 +3,7 @@ Utility
 
 .. automodule:: gdeep.utility
    :members:
-   
+
 Persistence Gradient
 --------------------
 
@@ -14,4 +14,4 @@ Extended Persistence
 --------------------
 
 .. automodule:: gdeep.utility.extended_persistence
-   :members:
\ No newline at end of file
+   :members:
diff --git a/docs/source/modules/visualisation.rst b/docs/source/modules/visualisation.rst
index cfa75c46..55beb8df 100644
--- a/docs/source/modules/visualisation.rst
+++ b/docs/source/modules/visualisation.rst
@@ -3,4 +3,3 @@ Visualisation
 
 .. automodule:: gdeep.visualisation
    :members:
-   
diff --git a/gdeep/analysis/decision_boundary/__init__.py b/gdeep/analysis/decision_boundary/__init__.py
index 5db47ece..68cc6b1e 100644
--- a/gdeep/analysis/decision_boundary/__init__.py
+++ b/gdeep/analysis/decision_boundary/__init__.py
@@ -1,14 +1,16 @@
 from .utility import UniformlySampledPoint
-from .decision_boundary_calculator import \
-    DecisionBoundaryCalculator,  \
-    GradientFlowDecisionBoundaryCalculator, \
-    QuasihyperbolicDecisionBoundaryCalculator
+from .decision_boundary_calculator import (
+    DecisionBoundaryCalculator,
+    GradientFlowDecisionBoundaryCalculator,
+    QuasihyperbolicDecisionBoundaryCalculator,
+)
+
 # from .compute_topology import
 
 
 __all__ = [
-    'UniformlySampledPoint',
-    'DecisionBoundaryCalculator',
-    'GradientFlowDecisionBoundaryCalculator',
-    'QuasihyperbolicDecisionBoundaryCalculator'
-    ]
+    "UniformlySampledPoint",
+    "DecisionBoundaryCalculator",
+    "GradientFlowDecisionBoundaryCalculator",
+    "QuasihyperbolicDecisionBoundaryCalculator",
+]
diff --git a/gdeep/analysis/decision_boundary/decision_boundary_calculator.py b/gdeep/analysis/decision_boundary/decision_boundary_calculator.py
index 4ee2a2e7..b638cf56 100644
--- a/gdeep/analysis/decision_boundary/decision_boundary_calculator.py
+++ b/gdeep/analysis/decision_boundary/decision_boundary_calculator.py
@@ -20,8 +20,7 @@ class DecisionBoundaryCalculator(ABC):
 
     @abstractmethod
     def step(self, number_of_steps=1):
-        """ Performs a single step towards the decision boundary
-        """
+        """Performs a single step towards the decision boundary"""
         pass
 
     @abstractmethod
@@ -49,6 +48,7 @@ def _convert_to_distance_to_boundary(self, model, output_shape):
         elif output_shape[-1] == 2:
             new_model = lambda x: torch.abs(model(x)[:, 0] - 0.5)
         else:
+
             def new_model(x):
                 y = torch.topk(model(x), 2).values
                 return y[:, 0] - y[:, 1]
@@ -72,9 +72,13 @@ class GradientFlowDecisionBoundaryCalculator(DecisionBoundaryCalculator):
             Function returning an optimizer for the params given as an
             argument.
     """
-    def __init__(self, model: Callable[[torch.Tensor], torch.Tensor],
-                 initial_points: torch.Tensor,
-                 optimizer: Callable[[list], torch.optim.Optimizer]):
+
+    def __init__(
+        self,
+        model: Callable[[torch.Tensor], torch.Tensor],
+        initial_points: torch.Tensor,
+        optimizer: Callable[[list], torch.optim.Optimizer],
+    ):
 
         self.sample_points = initial_points.to(DEVICE)
         self.sample_points.requires_grad = True
@@ -82,30 +86,29 @@ def __init__(self, model: Callable[[torch.Tensor], torch.Tensor],
         output_shape = output.size()
 
         if not len(output_shape) in [1, 2]:
-            raise RuntimeError('Output of model has wrong size!')
+            raise RuntimeError("Output of model has wrong size!")
         # Convert `model` to `self.model` with squared generalized distance
         # to the decision boundary as output.
         new_model = self._convert_to_distance_to_boundary(model, output_shape)
-        self.model = lambda x: new_model(x)**2
+        self.model = lambda x: new_model(x) ** 2
         # Check if self.model has the right output shape
-        assert len(self.model(self.sample_points).size()) == 1, \
-            f'Output shape is {self.model(self.sample_points).size()}'
+        assert (
+            len(self.model(self.sample_points).size()) == 1
+        ), f"Output shape is {self.model(self.sample_points).size()}"
 
         self.optimizer = optimizer([self.sample_points])
 
     def step(self, number_of_steps=1):
-        """Performs the indicated number of steps towards the decision boundary
-        """
+        """Performs the indicated number of steps towards the decision boundary"""
         print("Executing the decision boundary computations:")
         for j in range(number_of_steps):
-            print("Step: " + str(j) + "/" + str(number_of_steps),
-                  end ='\r')
+            print("Step: " + str(j) + "/" + str(number_of_steps), end="\r")
             self.optimizer.zero_grad()
             loss = torch.sum(self.model(self.sample_points)).to(DEVICE)
             loss.backward()
-            #if DEVICE.type == "xla":
+            # if DEVICE.type == "xla":
             #    xm.optimizer_step(self.optimizer, barrier=True)  # Note: Cloud TPU-specific code!
-            #else:
+            # else:
             self.optimizer.step()
 
     def get_decision_boundary(self) -> torch.Tensor:
@@ -120,10 +123,14 @@ class QuasihyperbolicDecisionBoundaryCalculator(DecisionBoundaryCalculator):
     """
     Computes Decision Boundary by emanating quasihyperbolic geodesics
     """
-    def __init__(self, model: Callable[[torch.Tensor], torch.Tensor],
-                 initial_points: torch.Tensor,
-                 initial_vectors: torch.Tensor,
-                 integrator=None):
+
+    def __init__(
+        self,
+        model: Callable[[torch.Tensor], torch.Tensor],
+        initial_points: torch.Tensor,
+        initial_vectors: torch.Tensor,
+        integrator=None,
+    ):
         """
         Args:
             model (Callable[[torch.Tensor], torch.Tensor]):
@@ -150,19 +157,18 @@ def __init__(self, model: Callable[[torch.Tensor], torch.Tensor],
         output_shape = output.size()
 
         if not len(output_shape) in [1, 2]:
-            raise RuntimeError('Output of model has wrong size!')
+            raise RuntimeError("Output of model has wrong size!")
         # Convert `model` to `self.model` with generalized distance
         # to the decision boundary as output.
-        self.model = self._convert_to_distance_to_boundary(model,
-                                                           output_shape)
+        self.model = self._convert_to_distance_to_boundary(model, output_shape)
 
         # Check if `self.model` has the right output shape
         distance_function = self.model(self.points)
-        assert len(distance_function.size()) == 1, \
-            f'Output shape is {distance_function.size()}'
+        assert (
+            len(distance_function.size()) == 1
+        ), f"Output shape is {distance_function.size()}"
         # Normalize tangent vectors in quasihyperbolic metric
-        self.vectors = initial_vectors * \
-            distance_function.unsqueeze(-1)
+        self.vectors = initial_vectors * distance_function.unsqueeze(-1)
 
         self.integrator = integrator
 
@@ -190,11 +196,9 @@ def odes(t, x):
             gradient_log_delta = gradient(y)
 
             # quasi-hyperbolic geodesic equation see markdown comment
-            ddy = 2*torch.einsum('bi,bi,bj->bj',
-                                 gradient_log_delta,
-                                 dy, dy) - \
-                torch.einsum('bi,bj,bj->bi', gradient_log_delta,
-                             dy, dy)
+            ddy = 2 * torch.einsum(
+                "bi,bi,bj->bj", gradient_log_delta, dy, dy
+            ) - torch.einsum("bi,bj,bj->bi", gradient_log_delta, dy, dy)
             return torch.stack((dy, ddy))
 
         ode_initial_conditions = torch.stack((self.points, self.vectors))
@@ -202,14 +206,17 @@ def odes(t, x):
         step_size = 1e-1
         # t = torch.arange(0, 3e0, 2e-3)
 
-        ode_solver_output = odeint(odes, ode_initial_conditions,
-                                   torch.tensor([0,
-                                                 number_of_steps
-                                                 * step_size],
-                                   dtype=self.points.dtype,
-                                   device=self.points.device),
-                                   method='rk4',
-                                   options=dict(step_size=step_size))
+        ode_solver_output = odeint(
+            odes,
+            ode_initial_conditions,
+            torch.tensor(
+                [0, number_of_steps * step_size],
+                dtype=self.points.dtype,
+                device=self.points.device,
+            ),
+            method="rk4",
+            options=dict(step_size=step_size),
+        )
         self.points = ode_solver_output[1, 0]
         self.vectors = ode_solver_output[1, 1]
 
diff --git a/gdeep/analysis/interpretability/__init__.py b/gdeep/analysis/interpretability/__init__.py
index 69a4231a..18e5185d 100644
--- a/gdeep/analysis/interpretability/__init__.py
+++ b/gdeep/analysis/interpretability/__init__.py
@@ -2,8 +2,4 @@
 from .attribution_factory import AttributionFactory, get_attr
 
 
-__all__ = [
-    'Interpreter',
-    'AttributionFactory',
-    'get_attr'
-    ]
+__all__ = ["Interpreter", "AttributionFactory", "get_attr"]
diff --git a/gdeep/analysis/interpretability/attribution_factory.py b/gdeep/analysis/interpretability/attribution_factory.py
index f2c62fe0..d349041f 100644
--- a/gdeep/analysis/interpretability/attribution_factory.py
+++ b/gdeep/analysis/interpretability/attribution_factory.py
@@ -1,4 +1,3 @@
-
 from typing import Dict, Any
 
 from captum import attr
@@ -9,7 +8,7 @@
 
 
 class AttributionFactory(object):
-    """ Attribution factory class for the Captum integration.
+    """Attribution factory class for the Captum integration.
     This factory will contain the attributions techniques of captum.
     """
 
@@ -17,13 +16,13 @@ class AttributionFactory(object):
 
     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
-        
+        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.
@@ -32,7 +31,7 @@ def build(self, key: str, *args, **kwargs) -> Any:  # type: ignore
             key:
                 the name of the dataset
         """
-        builder = self._builders.get(key)    # type: ignore
+        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
@@ -40,7 +39,7 @@ def build(self, key: str, *args, **kwargs) -> Any:  # type: ignore
 
 class AttributionBuilder(object):
     """Builder class for the torchvision dataset
-        and all its variations"""
+    and all its variations"""
 
     def __init__(self, attr_name: str):
         self.attr_name = attr_name
diff --git a/gdeep/analysis/interpretability/captum.py b/gdeep/analysis/interpretability/captum.py
index 13f25902..61173bed 100644
--- a/gdeep/analysis/interpretability/captum.py
+++ b/gdeep/analysis/interpretability/captum.py
@@ -25,6 +25,7 @@ class Interpreter:
             more info at https://captum.ai/tutorials/
 
     """
+
     x: Tensor
     attribution: Tensor
     sentence: str
@@ -33,18 +34,19 @@ class Interpreter:
     features_list: Optional[List]
     attribution_list: Optional[List]
 
-    def __init__(self, model: nn.Module,
-                 method: str = "IntegratedGradients"):
+    def __init__(self, model: nn.Module, method: str = "IntegratedGradients"):
         # self.model = model
         self.model = model.to(DEVICE)
         self.method = method
         self.stored_visualisations: List = []
 
-    def interpret(self,
-                  x: Tensor,
-                  y: Optional[Any] = None,
-                  layer: Optional[torch.nn.Module] = None,
-                  **kwargs: Any) -> Tuple[Tensor, Tensor]:
+    def interpret(
+        self,
+        x: Tensor,
+        y: Optional[Any] = None,
+        layer: Optional[torch.nn.Module] = None,
+        **kwargs: Any
+    ) -> Tuple[Tensor, Tensor]:
         """This method creates a datum interpreter. This class
         is based on captum.
 
@@ -78,8 +80,9 @@ def interpret(self,
             self.attribution = attr_class.attribute(self.x, **kwargs)
         return self.x, self.attribution
 
-    def feature_importance(self, x: Tensor, y: Tensor, **kwargs: Any) \
-            -> Tuple[Tensor, List[Any]]:
+    def feature_importance(
+        self, x: Tensor, y: Tensor, **kwargs: Any
+    ) -> Tuple[Tensor, List[Any]]:
         """This method creates a tabular data interpreter. This class
         is based on captum.
 
@@ -100,27 +103,42 @@ def feature_importance(self, x: Tensor, y: Tensor, **kwargs: Any) \
         self.y = y.to(DEVICE)
         self.model.eval()
         attr_list = []
-        attribute_dict = {"inputs": self.x, "target": self.y,
-                          "sliding_window_shapes": (1, ),
-                          "n_steps": 50, **kwargs
-                          }
-        self.features_list = ["IntegratedGradients", "IntegratedGradients",
-                              "DeepLift", "FeatureAblation", "Occlusion"]
+        attribute_dict = {
+            "inputs": self.x,
+            "target": self.y,
+            "sliding_window_shapes": (1,),
+            "n_steps": 50,
+            **kwargs,
+        }
+        self.features_list = [
+            "IntegratedGradients",
+            "IntegratedGradients",
+            "DeepLift",
+            "FeatureAblation",
+            "Occlusion",
+        ]
         for method in self.features_list:
             attribution_class = get_attr(method, self.model)
-            args_names = inspect.signature(attribution_class.attribute).parameters.keys()
-            kwargs_of_method = {k: v for k, v in attribute_dict.items() if k in args_names}
+            args_names = inspect.signature(
+                attribution_class.attribute
+            ).parameters.keys()
+            kwargs_of_method = {
+                k: v for k, v in attribute_dict.items() if k in args_names
+            }
             attr_list.append(attribution_class.attribute(**kwargs_of_method))
         self.attribution_list = attr_list
         return self.x, self.attribution_list
 
-    def interpret_text(self, sentence: str,
-                       label: Any,
-                       vocab: Dict[str, int],
-                       tokenizer: Callable[[str], List[str]],
-                       layer: Optional[torch.nn.Module] = None,
-                       min_len: int = 7,
-                       **kwargs) -> Tuple[str, Tensor]:
+    def interpret_text(
+        self,
+        sentence: str,
+        label: Any,
+        vocab: Dict[str, int],
+        tokenizer: Callable[[str], List[str]],
+        layer: Optional[torch.nn.Module] = None,
+        min_len: int = 7,
+        **kwargs
+    ) -> Tuple[str, Tensor]:
         """This method creates a text interpreter. This class
         is based on captum.
 
@@ -153,7 +171,7 @@ def interpret_text(self, sentence: str,
             attr_class = get_attr(self.method, self.model)
         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()
@@ -171,43 +189,54 @@ def interpret_text(self, sentence: str,
         # generate reference indices for each sample
         pad_index = 0
         token_reference = TokenReferenceBase(reference_token_idx=pad_index)
-        reference_indices = \
-            token_reference.generate_reference(seq_length,
-                                               DEVICE).unsqueeze(0)
+        reference_indices = token_reference.generate_reference(
+            seq_length, DEVICE
+        ).unsqueeze(0)
         # compute attributions and approximation
         # delta using layer integrated gradients
-        self.attribution, delta = attr_class.attribute(input_indices,
-                                                       reference_indices,
-                                                       target=label,
-                                                       **kwargs)
-
-        print('pred: ', pred_ind, '(', '%.2f' % pred, ')',
-              ', delta: ', abs(delta.item()))
-
-        self.add_attributions_to_visualizer(self.attribution, text, pred,
-                                            pred_ind, label, delta.item(),
-                                            self.stored_visualisations)
+        self.attribution, delta = attr_class.attribute(
+            input_indices, reference_indices, target=label, **kwargs
+        )
+
+        print(
+            "pred: ", pred_ind, "(", "%.2f" % pred, ")", ", delta: ", abs(delta.item())
+        )
+
+        self.add_attributions_to_visualizer(
+            self.attribution,
+            text,
+            pred,
+            pred_ind,
+            label,
+            delta.item(),
+            self.stored_visualisations,
+        )
         return sentence, self.attribution
 
     @staticmethod
-    def add_attributions_to_visualizer(attributions,
-                                       text: List[str],
-                                       pred: Tensor,
-                                       pred_ind: float,
-                                       label: Any,
-                                       delta: Any,
-                                       vis_data_records: List) -> None:
+    def add_attributions_to_visualizer(
+        attributions,
+        text: List[str],
+        pred: Tensor,
+        pred_ind: float,
+        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()
 
         # storing couple samples in an array for visualization purposes
-        vis_data_records.append(visualization.VisualizationDataRecord(
-                                attributions,
-                                pred,
-                                str(pred_ind),
-                                str(label),
-                                "1",
-                                attributions.sum(),
-                                text,
-                                delta))
+        vis_data_records.append(
+            visualization.VisualizationDataRecord(
+                attributions,
+                pred,
+                str(pred_ind),
+                str(label),
+                "1",
+                attributions.sum(),
+                text,
+                delta,
+            )
+        )
diff --git a/gdeep/data/__init__.py b/gdeep/data/__init__.py
index 913e90a4..4d4a8221 100644
--- a/gdeep/data/__init__.py
+++ b/gdeep/data/__init__.py
@@ -13,5 +13,5 @@
     "MissingVocabularyError",
     "datasets",
     "persistence_diagrams",
-    "preprocessors"
+    "preprocessors",
 ]
diff --git a/gdeep/data/dataset_factory.py b/gdeep/data/dataset_factory.py
index b413197a..3e5452fe 100644
--- a/gdeep/data/dataset_factory.py
+++ b/gdeep/data/dataset_factory.py
@@ -7,14 +7,14 @@
 
 
 class DatasetFactory(object):
-    """ Dataset factory class for the tori dataset and torchvision datasets
+    """Dataset factory class for the tori dataset and torchvision datasets
     using the factory design pattern
-    
+
     Examples::
 
         # Create a dataset for the tori dataset
         dataset = get_dataset("Tori", name="DoubleTori", n_points=100)
-        
+
         # Create the MNIST dataset
         dataset = get_dataset("Torchvision", name="MNIST")
 
diff --git a/gdeep/data/datasets/__init__.py b/gdeep/data/datasets/__init__.py
index d1849223..8d5a859e 100644
--- a/gdeep/data/datasets/__init__.py
+++ b/gdeep/data/datasets/__init__.py
@@ -3,31 +3,34 @@
 from .dataset_cloud import DatasetCloud
 from ._data_cloud import _DataCloud
 from .build_datasets import DatasetBuilder, get_dataset
-from .base_dataloaders import DataLoaderBuilder, \
-    AbstractDataLoaderBuilder
+from .base_dataloaders import DataLoaderBuilder, AbstractDataLoaderBuilder
 from .dataset_for_image import ImageClassificationFromFiles
 from .dataset_form_array import FromArray
 from .dataloader_cloud import DlBuilderFromDataCloud
-from .parallel_orbit import generate_orbit_parallel, create_pd_orbits,\
-    OrbitsGenerator, DataLoaderKwargs
+from .parallel_orbit import (
+    generate_orbit_parallel,
+    create_pd_orbits,
+    OrbitsGenerator,
+    DataLoaderKwargs,
+)
 from .persistence_diagrams_from_files import PersistenceDiagramFromFiles
 
 
 __all__ = [
-    'Rotation',
-    'CategoricalDataCloud',
-    'AbstractDataLoaderBuilder',
-    'ToriDataset',
-    'DataLoaderBuilder',
-    'DatasetBuilder',
-    'generate_orbit_parallel',
-    'create_pd_orbits',
-    'OrbitsGenerator',
-    'DataLoaderKwargs',
-    'get_dataset',
-    'ImageClassificationFromFiles',
-    'FromArray',
-    'DatasetCloud',
-    'DlBuilderFromDataCloud',
-    'PersistenceDiagramFromFiles'
-    ]
+    "Rotation",
+    "CategoricalDataCloud",
+    "AbstractDataLoaderBuilder",
+    "ToriDataset",
+    "DataLoaderBuilder",
+    "DatasetBuilder",
+    "generate_orbit_parallel",
+    "create_pd_orbits",
+    "OrbitsGenerator",
+    "DataLoaderKwargs",
+    "get_dataset",
+    "ImageClassificationFromFiles",
+    "FromArray",
+    "DatasetCloud",
+    "DlBuilderFromDataCloud",
+    "PersistenceDiagramFromFiles",
+]
diff --git a/gdeep/data/datasets/_data_cloud.py b/gdeep/data/datasets/_data_cloud.py
index 7efd84d3..01cc86ec 100644
--- a/gdeep/data/datasets/_data_cloud.py
+++ b/gdeep/data/datasets/_data_cloud.py
@@ -18,17 +18,22 @@
 LOGGER = logging.getLogger(__name__)
 LOGGER.setLevel(logging.WARNING)
 
+
 def _check_public_access(use_public_access: bool):
     """Check if the public access is enabled."""
+
     def wrap(function):
         def wrapper_function(*args, **kwargs):
-            if(use_public_access):
+            if use_public_access:
                 raise ValueError("DataCloud object has public access only!")
             return function(*args, **kwargs)
+
         return wrapper_function
+
     return wrap
 
-class _DataCloud():
+
+class _DataCloud:
     """Download handle for Google Cloud Storage buckets.
 
     Args:
@@ -45,20 +50,21 @@ class _DataCloud():
             Path to the credentials file.
             Only used if public_access is False and credentials are not
             provided. Defaults to None.
-            
+
         Raises:
             ValueError: If the bucket does not exist.
-            
+
         Returns:
             None
     """
+
     def __init__(
-            self,
-            bucket_name: str = DATASET_BUCKET_NAME,
-            download_directory: str = DEFAULT_DOWNLOAD_DIR,
-            use_public_access: bool = True,
-            path_to_credentials: Union[str, None] = None,
-            ) -> None:
+        self,
+        bucket_name: str = DATASET_BUCKET_NAME,
+        download_directory: str = DEFAULT_DOWNLOAD_DIR,
+        use_public_access: bool = True,
+        path_to_credentials: Union[str, None] = None,
+    ) -> None:
         self.bucket_name = bucket_name
         self.use_public_access = use_public_access
         if not self.use_public_access:
@@ -66,45 +72,44 @@ def __init__(
             if path_to_credentials is None:
                 self.storage_client = storage.Client()
             else:
-                credentials = service_account.Credentials.\
-                    from_service_account_file(path_to_credentials)
+                credentials = service_account.Credentials.from_service_account_file(
+                    path_to_credentials
+                )
                 self.storage_client = storage.Client(credentials=credentials)
             self.bucket = self.storage_client.bucket(self.bucket_name)
         else:
-            self.public_url = ("https://storage.googleapis.com/" + 
-                               bucket_name + "/")
-        
+            self.public_url = "https://storage.googleapis.com/" + bucket_name + "/"
+
         # Set up download path
         self.download_directory = download_directory
-        
-        # Create a new directory because it does not exist 
-        if not exists(self.download_directory) \
-            and self.download_directory != "":
-              makedirs(self.download_directory)
+
+        # Create a new directory because it does not exist
+        if not exists(self.download_directory) and self.download_directory != "":
+            makedirs(self.download_directory)
 
     def list_blobs(self) -> List[str]:
         """List all blobs in the bucket.
-        
+
         Returns:
             List[str]:
                 List of blobs in the bucket.
         """
         # Assert that the bucket does not use public access
         if self.use_public_access:
-            raise ValueError("DataCloud object can only list blobs"
-                             "when using private access!")
+            raise ValueError(
+                "DataCloud object can only list blobs" "when using private access!"
+            )
         blobs = self.bucket.list_blobs()
         return [blob.name for blob in blobs]
-    
-    def blob_exists(self,
-                    blob_name: str) -> bool:
+
+    def blob_exists(self, blob_name: str) -> bool:
         """Check if a Blob exists in the bucket.
 
         Args:
             blob_name (str):
                 Name of the Blob to check.
-                
-            
+
+
         Returns:
             bool:
                 True if the Blob exists, False otherwise.
@@ -115,11 +120,11 @@ def blob_exists(self,
             return response.status_code == 200
         else:
             blob = self.bucket.blob(blob_name)
-            return blob.exists() 
-        
-    def download_file(self,
-                 blob_name: str,
-                 download_directory: Union[str, None]=None) -> None:
+            return blob.exists()
+
+    def download_file(
+        self, blob_name: str, download_directory: Union[str, None] = None
+    ) -> None:
         """Download a blob from Google Cloud Storage bucket.
 
         Args:
@@ -128,11 +133,11 @@ def download_file(self,
                 root of the bucket.
             download_directory (str, optional):
                 Directory to download the blob to.
-        
+
         Raises:
             ValueError:
                 If the blob does not exist.
-            
+
         Returns:
             None
         """
@@ -145,7 +150,7 @@ def download_file(self,
             raise google.api_core.exceptions.NotFound(
                 "Blob {} does not exist!".format(blob_name)
             )
-        
+
         # If the file exists, compare checksums
         if isfile(join(download_directory, blob_name)):
             # Get remote md5 checksum from url in base64 format.
@@ -157,60 +162,62 @@ def download_file(self,
             else:
                 blob = self.bucket.blob(blob_name)
                 checksum_remote = blob.md5_hash
-                
+
             checksum_local = get_checksum(
-                join(download_directory, blob_name),
-                encoding="base64"
-                )
+                join(download_directory, blob_name), encoding="base64"
+            )
             if checksum_remote is not None:
                 if checksum_remote != checksum_local:
                     # Ask user if they want to download the file
                     answer = input(
-                        f"File {join(download_directory, blob_name)} already" +
-                        "exists and checksums don't match! " +
-                        "Do you want to overwrite it? [y/N]")
-                    if answer.lower() not in  ["y", "yes"]:
+                        f"File {join(download_directory, blob_name)} already"
+                        + "exists and checksums don't match! "
+                        + "Do you want to overwrite it? [y/N]"
+                    )
+                    if answer.lower() not in ["y", "yes"]:
                         return
-                else: 
-                    print(f"File {join(download_directory, blob_name)} " +
-                          "already exists and checksums match! " +
-                                "Skipping download.")
+                else:
+                    print(
+                        f"File {join(download_directory, blob_name)} "
+                        + "already exists and checksums match! "
+                        + "Skipping download."
+                    )
                     return
             else:
-                print(f"File {join(download_directory, blob_name)} already" +
-                      "exists and remote checksum is " +
-                            "None! Downloading anyway.")
+                print(
+                    f"File {join(download_directory, blob_name)} already"
+                    + "exists and remote checksum is "
+                    + "None! Downloading anyway."
+                )
         # Download file
-        print("Downloading file {} to {}".format(
-            blob_name, download_directory))
+        print("Downloading file {} to {}".format(blob_name, download_directory))
         if self.use_public_access:
             wget.download(url, join(download_directory, blob_name))
         else:
-            self.bucket.blob(blob_name)\
-                .download_to_filename(join(download_directory, blob_name),
-                                      checksum="md5")
-        
-    def download_folder(self,
-                        blob_name: str) -> None:
+            self.bucket.blob(blob_name).download_to_filename(
+                join(download_directory, blob_name), checksum="md5"
+            )
+
+    def download_folder(self, blob_name: str) -> None:
         """Download a folder from Google Cloud Storage bucket.
-        
+
         Warning: This function does not download empty subdirectories.
 
         Args:
             blob_name (str):
                 Name of the blob folder to download. The name is relative to
                 the root of the bucket.
-        
+
         Raises:
             RuntimeError:
                 If the folder does not exist.
-        
+
         Returns:
             None
         """
-        assert not self.use_public_access, ("Downloading folders is not"
-                                            "is not supported with public"
-                                            "access!")
+        assert not self.use_public_access, (
+            "Downloading folders is not" "is not supported with public" "access!"
+        )
         # Get list of files in the blob
         blobs = self.bucket.list_blobs(prefix=blob_name)
         for blob in blobs:
@@ -220,21 +227,24 @@ def download_folder(self,
             file_split = blob.name.split("/")
             directory = "/".join(file_split[0:-1])
             if not exists(directory):
-                makedirs(join(self.download_directory, directory),
-                         exist_ok=True)
+                makedirs(join(self.download_directory, directory), exist_ok=True)
             logging.getLogger().info("Downloading blob %s", blob.name)
-            
-            local_path =  blob.name.replace("/", "\\") \
-                if sys.platform == 'win32' else blob.name
-            
-            blob.download_to_filename(join(self.download_directory,local_path),
-                                      checksum="md5")
-        
-    def upload_file(self,
-               source_file_name: str,
-               target_blob_name: Union[str, None] = None,
-               make_public: bool = False,
-               overwrite: bool = False) -> None:
+
+            local_path = (
+                blob.name.replace("/", "\\") if sys.platform == "win32" else blob.name
+            )
+
+            blob.download_to_filename(
+                join(self.download_directory, local_path), checksum="md5"
+            )
+
+    def upload_file(
+        self,
+        source_file_name: str,
+        target_blob_name: Union[str, None] = None,
+        make_public: bool = False,
+        overwrite: bool = False,
+    ) -> None:
         """Upload a local file to Google Cloud Storage bucket.
 
         Args:
@@ -250,10 +260,10 @@ def upload_file(self,
                 overwrite (bool, optional):
             Whether or not to overwrite the target
                 Blob. Defaults to False.
-            
+
         Raises:
             RuntimeError: If the target Blob already exists.
-            
+
         Returns:
             None
         """
@@ -264,24 +274,24 @@ def upload_file(self,
             raise RuntimeError(f"Blob {target_blob_name} already exists.")
         logging.getLogger().info("upload file %s", source_file_name)
         # Check if source_file_name is bigger than 5GB
-        if isfile(source_file_name) and\
-            getsize(source_file_name) > 5000000000:
+        if isfile(source_file_name) and getsize(source_file_name) > 5000000000:
             raise ValueError("File is bigger than 5GB")
-        
+
         # Compute MD5 checksum of the file and add it to the metadata of
         # the blob
         blob.md5_hash = get_checksum(source_file_name, encoding="base64")
-        
+
         blob.upload_from_filename(source_file_name, checksum="md5")
         if make_public:
             blob.make_public()
-    
-    def upload_folder(self,
-                      source_folder: str,
-                      target_folder: str = None,
-                      make_public: bool = False,
-                      ) -> None:
-        """Upload a local folder with all it's subolders to Google 
+
+    def upload_folder(
+        self,
+        source_folder: str,
+        target_folder: str = None,
+        make_public: bool = False,
+    ) -> None:
+        """Upload a local folder with all it's subolders to Google
         Cloud Storage bucket.
 
         Args:
@@ -294,90 +304,95 @@ def upload_folder(self,
             make_public (bool, optional):
                 Whether or not to make the uploaded
                 file public. Defaults to False.
-        
+
         Raises:
             ValueError:
                 If the source folder is not a directory.
-            
+
         Returns:
             None
         """
         if not isdir(source_folder):
             raise ValueError("Source folder is not a directory.")
-        
+
         if target_folder is None:
             target_folder = ""
-        
+
         # List of all files in the source folder
-        files = [join(source_folder, f) for f in listdir(source_folder)
-                 if isfile(join(source_folder, f))]
+        files = [
+            join(source_folder, f)
+            for f in listdir(source_folder)
+            if isfile(join(source_folder, f))
+        ]
         print(files)
         # Upload all files in the source folder
         for file in files:
             file_name = os.path.basename(file)
-            if(target_folder == ""):
-                self.upload_file(join(source_folder, file_name),
-                                target_blob_name=file_name,
-                                make_public=make_public)
+            if target_folder == "":
+                self.upload_file(
+                    join(source_folder, file_name),
+                    target_blob_name=file_name,
+                    make_public=make_public,
+                )
             else:
-                self.upload_file(join(source_folder, file_name),
-                                target_blob_name=target_folder + '/' + file_name,
-                                make_public=make_public)
-        
+                self.upload_file(
+                    join(source_folder, file_name),
+                    target_blob_name=target_folder + "/" + file_name,
+                    make_public=make_public,
+                )
+
         # List of all subfolders in the source folder
-        subfolders = [join(source_folder, f) for f in listdir(source_folder)
-                        if isdir(join(source_folder, f))]
+        subfolders = [
+            join(source_folder, f)
+            for f in listdir(source_folder)
+            if isdir(join(source_folder, f))
+        ]
         # Upload all subfolders in the source folder recursively
         for subfolder in subfolders:
             relative_subfolder = os.path.relpath(subfolder, source_folder)
             if target_folder == "":
-                self.upload_folder(join(source_folder, relative_subfolder),
-                                target_folder=relative_subfolder,
-                                make_public=make_public)
+                self.upload_folder(
+                    join(source_folder, relative_subfolder),
+                    target_folder=relative_subfolder,
+                    make_public=make_public,
+                )
             else:
-                self.upload_folder(join(source_folder, relative_subfolder),
-                                target_folder=target_folder + '/' +\
-                                    relative_subfolder,
-                                make_public=make_public)
-
-        
-    
-    def delete_blob(self,
-                    blob_name: str) -> None:
+                self.upload_folder(
+                    join(source_folder, relative_subfolder),
+                    target_folder=target_folder + "/" + relative_subfolder,
+                    make_public=make_public,
+                )
+
+    def delete_blob(self, blob_name: str) -> None:
         """Deletes a single Blob from Google Cloud Storage
 
         Args:
             blob_name (str):
                 The name of the Blob to delete
-                
+
         Raises:
             RuntimeError: If the Blob does not exist.
-        
+
         Returns:
             None
         """
         blob = self.bucket.blob(blob_name)
         blob.delete()
-        
-    def delete_blobs(self,
-                     blobs_name: str) -> None:
+
+    def delete_blobs(self, blobs_name: str) -> None:
         """Deletes a Blob and all its children from Google Cloud Storage.
 
         Args:
             blobs_name (str):
                 Name of the parent Blob to delete.
-                
+
         Raises:
-            ValueError: 
+            ValueError:
                 If the Blob does not exist.
-            
+
         Returns:
             None
         """
         blobs = self.bucket.list_blobs(prefix=blobs_name)
         for blob in blobs:
             blob.delete()
-
-
-
-
diff --git a/gdeep/data/datasets/base_dataloaders.py b/gdeep/data/datasets/base_dataloaders.py
index a10ded12..735ff1f7 100644
--- a/gdeep/data/datasets/base_dataloaders.py
+++ b/gdeep/data/datasets/base_dataloaders.py
@@ -6,8 +6,18 @@
 from abc import ABC, abstractmethod
 from os.path import join
 from collections.abc import Iterable
-from typing import Any, Callable, Dict, Generic, Optional, Sequence, Tuple, \
-    TypeVar, Union, List
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+    List,
+)
 
 import torch
 from torch.utils.data import DataLoader, Dataset
@@ -21,16 +31,18 @@
 
 
 Tensor = torch.Tensor
-T = TypeVar('T')
+T = TypeVar("T")
 
 
 class AbstractDataLoaderBuilder(ABC):
     """The abstractr class to interface the
     Giotto dataloaders"""
+
     @abstractmethod
     def build(self, tuple_of_kwargs: List[Dict[str, Any]]):
         pass
 
+
 @dataclass
 class DataLoaderParams:
     batch_size: Optional[int] = 1
@@ -41,30 +53,29 @@ class DataLoaderParams:
     pin_memory: bool = False
     drop_last: bool = False
     timeout: float = 0.0
-    persistent_workers:bool = False
-    
+    persistent_workers: bool = False
+
     def copy(self):
         return DataLoaderParams(**self.to_dict())
-    
+
     def update_batch_size(self, batch_size: int):
         self.batch_size = batch_size
         return self
-    
+
     def update_shuffle(self, shuffle: bool):
         self.shuffle = shuffle
         return self
-    
+
     def to_dict(self):
-        return {
-            k: v for k, v in self.__dict__.items() if not k.startswith("_")
-        }
+        return {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
+
 
 @dataclass
 class DataLoaderParamsTuples:
     train: DataLoaderParams
     test: Optional[DataLoaderParams] = None
     validation: Optional[DataLoaderParams] = None
-        
+
     @staticmethod
     def default(
         collate_fn: Callable[[Any], Any],
@@ -73,7 +84,7 @@ def default(
         pin_memory: bool = False,
         drop_last: bool = False,
         timeout: float = 0.0,
-        persistent_workers:bool = False,
+        persistent_workers: bool = False,
         with_validation: bool = False,
     ) -> "DataLoaderParamsTuples":
         dlp = DataLoaderParamsTuples(
@@ -85,7 +96,7 @@ def default(
                 pin_memory=pin_memory,
                 drop_last=drop_last,
                 timeout=timeout,
-                persistent_workers=persistent_workers
+                persistent_workers=persistent_workers,
             ),
             test=DataLoaderParams(
                 batch_size=batch_size,
@@ -95,8 +106,8 @@ def default(
                 pin_memory=pin_memory,
                 drop_last=drop_last,
                 timeout=timeout,
-                persistent_workers=persistent_workers
-            )
+                persistent_workers=persistent_workers,
+            ),
         )
         if with_validation:
             dlp.validation = DataLoaderParams(
@@ -107,31 +118,26 @@ def default(
                 pin_memory=pin_memory,
                 drop_last=drop_last,
                 timeout=timeout,
-                persistent_workers=persistent_workers
+                persistent_workers=persistent_workers,
             )
         return dlp
-        
+
     def to_tuple_of_dicts(self) -> Tuple[Dict, ...]:
         if self.validation is not None and self.test is not None:
             return (
                 self.train.to_dict(),
                 self.validation.to_dict(),
-                self.test.to_dict()
+                self.test.to_dict(),
             )
         elif self.test is not None:
-            return (
-                self.train.to_dict(),
-                self.test.to_dict()
-            )
+            return (self.train.to_dict(), self.test.to_dict())
         else:
-            return (
-                self.train.to_dict(),
-            )
-        
+            return (self.train.to_dict(),)
+
     @staticmethod
     def from_list_of_dicts(
         list_of_kwargs: List[Dict[str, Any]]
-    ) -> 'DataLoaderParamsTuples':
+    ) -> "DataLoaderParamsTuples":
         """This method accepts the arguments of the torch
         Dataloader and applies them when creating the
         tuple
@@ -145,24 +151,21 @@ def from_list_of_dicts(
             return DataLoaderParamsTuples(
                 DataLoaderParams(**list_of_kwargs[0]),
                 DataLoaderParams(**list_of_kwargs[1]),
-                DataLoaderParams(**list_of_kwargs[2])
+                DataLoaderParams(**list_of_kwargs[2]),
             )
         elif len(list_of_kwargs) == 2:
             return DataLoaderParamsTuples(
                 DataLoaderParams(**list_of_kwargs[0]),
                 DataLoaderParams(**list_of_kwargs[1]),
-                None
+                None,
             )
         elif len(list_of_kwargs) == 1:
             return DataLoaderParamsTuples(
-                DataLoaderParams(**list_of_kwargs[0]),
-                None,
-                None
+                DataLoaderParams(**list_of_kwargs[0]), None, None
             )
         else:
-            raise ValueError(
-                "The list of dictionaries must have 2 or 3 elements"
-            )
+            raise ValueError("The list of dictionaries must have 2 or 3 elements")
+
 
 class DataLoaderBuilder(AbstractDataLoaderBuilder):
     """This class builds, out of a tuple of datasets, the
@@ -176,13 +179,17 @@ class DataLoaderBuilder(AbstractDataLoaderBuilder):
             they will be considered as training first and
             validation afterwards.
     """
+
     def __init__(self, tuple_of_datasets: List[Dataset[Any]]) -> None:
         self.tuple_of_datasets = tuple_of_datasets
         assert len(tuple_of_datasets) <= 3, "Too many Dataset inserted: maximum 3."
 
-    def build(self,
-              tuple_of_kwargs:Union[List[Dict[str, Any]], DataLoaderParamsTuples, None]=None
-              ) -> List[DataLoader[Any]]:
+    def build(
+        self,
+        tuple_of_kwargs: Union[
+            List[Dict[str, Any]], DataLoaderParamsTuples, None
+        ] = None,
+    ) -> List[DataLoader[Any]]:
         """This method accepts the arguments of the torch
         Dataloader and applies them when creating the
         tuple
@@ -192,23 +199,23 @@ def build(self,
                 List of dictionaries, each one being the
                 kwargs for the corresponding DataLoader
         """
-        out: List = []  
+        out: List = []
         if tuple_of_kwargs is None:
             for i, dataset in enumerate(self.tuple_of_datasets):
                 out.append(DataLoader(dataset))
             out += [None] * (3 - len(out))
             return out
-        
+
         if isinstance(tuple_of_kwargs, DataLoaderParamsTuples):
             tuple_of_kwargs = tuple_of_kwargs.to_tuple_of_dicts()  # type: ignore
-        
-        assert isinstance(tuple_of_kwargs, (list, tuple)), ("The kwargs must be a list or a tuple at"
-            "this point")
-        assert len(tuple_of_kwargs) == len(self.tuple_of_datasets), \
-            "Cannot match the dataloaders and the parameters. "
+
+        assert isinstance(tuple_of_kwargs, (list, tuple)), (
+            "The kwargs must be a list or a tuple at" "this point"
+        )
+        assert len(tuple_of_kwargs) == len(
+            self.tuple_of_datasets
+        ), "Cannot match the dataloaders and the parameters. "
         for dataset, kwargs in zip(self.tuple_of_datasets, tuple_of_kwargs):
             out.append(DataLoader(dataset, **kwargs))
         out += [None] * (3 - len(out))
         return out
-
-
diff --git a/gdeep/data/datasets/build_datasets.py b/gdeep/data/datasets/build_datasets.py
index 4cfdc971..b855b585 100644
--- a/gdeep/data/datasets/build_datasets.py
+++ b/gdeep/data/datasets/build_datasets.py
@@ -10,12 +10,12 @@
 from ..dataset_factory import DatasetFactory
 
 Tensor = torch.Tensor
-T = TypeVar('T')
+T = TypeVar("T")
 
 
 class TorchvisionDatasetBuilder(object):
     """Builder class for the torchvision dataset
-        and all its variations"""
+    and all its variations"""
 
     def __init__(self, dataset_name: str):
         self.dataset_name = dataset_name
@@ -28,7 +28,7 @@ def __call__(self, **kwargs):  # type: ignore
 
 class TorchtextDatasetBuilder(object):
     """Builder class for the torchtext dataset
-        and all its variations"""
+    and all its variations"""
 
     def __init__(self, dataset_name: str):
         self.dataset_name = dataset_name
@@ -51,7 +51,7 @@ def __call__(self, **kwargs):  # type: ignore
 
 
 def get_dataset(key: str, **kwargs) -> Tuple[Dataset[Any]]:
-    """ Get a dataset from the factory
+    """Get a dataset from the factory
 
     Args:
         key :
@@ -94,17 +94,20 @@ class DatasetBuilder:
             IterableDataset
 
     """
+
     train_ds: Dataset[Any]
     valid_ds: Optional[Dataset[Any]]
     test_ds: Optional[Dataset[Any]]
 
-    def __init__(self, name: str = "MNIST", convert_to_map_dataset: bool = False) -> None:
+    def __init__(
+        self, name: str = "MNIST", convert_to_map_dataset: bool = False
+    ) -> None:
         self.convert_to_map_dataset = convert_to_map_dataset
         self.name = name
 
-    def build(self, **kwargs) -> Tuple[Dataset[Any],
-                                       Optional[Dataset[Any]],
-                                       Optional[Dataset[Any]]]:
+    def build(
+        self, **kwargs
+    ) -> Tuple[Dataset[Any], Optional[Dataset[Any]], Optional[Dataset[Any]]]:
         """Method that returns the dataset.
 
         Args:
@@ -133,16 +136,17 @@ def build(self, **kwargs) -> Tuple[Dataset[Any],
             self.test_ds = None
 
         if self.convert_to_map_dataset:
-            self.train_ds, self.valid_ds, self.test_ds = self._convert(self.train_ds,
-                                                                       self.valid_ds,
-                                                                       self.test_ds)
+            self.train_ds, self.valid_ds, self.test_ds = self._convert(
+                self.train_ds, self.valid_ds, self.test_ds
+            )
         return self.train_ds, self.valid_ds, self.test_ds
 
-    def _convert(self, training_data: Dataset[Any],
-                 validation_data: Optional[Dataset[Any]] = None,
-                 test_data: Optional[Dataset[Any]] = None) -> Tuple[Dataset[Any],
-                                                                    Optional[Dataset[Any]],
-                                                                    Optional[Dataset[Any]]]:
+    def _convert(
+        self,
+        training_data: Dataset[Any],
+        validation_data: Optional[Dataset[Any]] = None,
+        test_data: Optional[Dataset[Any]] = None,
+    ) -> Tuple[Dataset[Any], Optional[Dataset[Any]], Optional[Dataset[Any]]]:
         """This private method converts and IterableDataset
         to a MapDataset"""
         if isinstance(training_data, torch.utils.data.IterableDataset):
@@ -156,7 +160,9 @@ def _convert(self, training_data: Dataset[Any],
         return training_data, validation_data, test_data
 
     @staticmethod
-    def _to_map_style_dataset(dataset: Union[Dataset, IterDataPipe]) -> Union[Dataset, MapDataPipe]:
+    def _to_map_style_dataset(
+        dataset: Union[Dataset, IterDataPipe]
+    ) -> Union[Dataset, MapDataPipe]:
         """our version of to_map_style_dataset, with the
         possibility to handle DataPipe and Dataset as well
         """
diff --git a/gdeep/data/datasets/categorical_data.py b/gdeep/data/datasets/categorical_data.py
index 78acb29e..30e0c817 100644
--- a/gdeep/data/datasets/categorical_data.py
+++ b/gdeep/data/datasets/categorical_data.py
@@ -1,5 +1,2 @@
-class CategoricalDataCloud():
+class CategoricalDataCloud:
     pass
-
-
-
diff --git a/gdeep/data/datasets/dataloader_cloud.py b/gdeep/data/datasets/dataloader_cloud.py
index 325fcb43..1c422d0a 100644
--- a/gdeep/data/datasets/dataloader_cloud.py
+++ b/gdeep/data/datasets/dataloader_cloud.py
@@ -14,12 +14,12 @@
 from .base_dataloaders import AbstractDataLoaderBuilder
 
 Tensor = torch.Tensor
-T = TypeVar('T')
+T = TypeVar("T")
 
 
 class DlBuilderFromDataCloud(AbstractDataLoaderBuilder):
     """Class that loads data from Google Cloud Storage
-    
+
     This class is useful to build dataloaders from a dataset stored in
     the GDeep Dataset Cloud on Google Cloud Storage.
 
@@ -28,7 +28,7 @@ class DlBuilderFromDataCloud(AbstractDataLoaderBuilder):
     to the download directory. The dataset is downloaded in the version
     used by Datasets Cloud, which may be different from the version
     used by the dataset's original developers.
-    
+
     Args:
         dataset_name (str):
             The name of the dataset.
@@ -42,7 +42,7 @@ class DlBuilderFromDataCloud(AbstractDataLoaderBuilder):
             Path to the credentials file.
             Only used if public_access is False and credentials are not
             provided. Defaults to None.
-        
+
 
     Returns:
         torch.utils.data.DataLoader: The dataloader for the dataset.
@@ -55,107 +55,137 @@ class DlBuilderFromDataCloud(AbstractDataLoaderBuilder):
             If the download_directory is not a valid directory.
     """
 
-    def __init__(self,
-                 dataset_name: str,
-                 download_directory: str,
-                 use_public_access: bool = True,
-                 path_to_credentials: Union[None, str] = None,
-                 ):
+    def __init__(
+        self,
+        dataset_name: str,
+        download_directory: str,
+        use_public_access: bool = True,
+        path_to_credentials: Union[None, str] = None,
+    ):
         self.dataset_name = dataset_name
         self.download_directory = download_directory
 
         # Download the dataset if it does not exist
         self.download_directory
 
-        self._download_dataset(use_public_access=use_public_access,
-                               path_to_credentials=path_to_credentials)
+        self._download_dataset(
+            use_public_access=use_public_access, path_to_credentials=path_to_credentials
+        )
 
         self.dl_builder = None
 
         # Load the metadata of the dataset
-        with open(join(self.download_directory, self.dataset_name,
-                       "metadata.json")) as f:
+        with open(
+            join(self.download_directory, self.dataset_name, "metadata.json")
+        ) as f:
             self.dataset_metadata = json.load(f)
 
         # Load the data and labels of the dataset
-        if self.dataset_metadata['data_type'] == 'tabular':
-            if self.dataset_metadata['data_format'] == 'pytorch_tensor':
-                data = torch.load(join(self.download_directory,
-                                       self.dataset_name, "data.pt"))
-                labels = torch.load(join(self.download_directory,
-                                         self.dataset_name, "labels.pt"))
-
-                self.dl_builder = DataLoaderBuilder([FromArray(data, labels), ])
-            elif self.dataset_metadata['data_format'] == 'numpy_array':
-                data = np.load(join(self.download_directory,
-                                    self.dataset_name, "data.npy"))
-                labels = np.load(join(self.download_directory,
-                                      self.dataset_name, "labels.npy"))
-                self.dl_builder = DataLoaderBuilder([FromArray(data, labels), ])
+        if self.dataset_metadata["data_type"] == "tabular":
+            if self.dataset_metadata["data_format"] == "pytorch_tensor":
+                data = torch.load(
+                    join(self.download_directory, self.dataset_name, "data.pt")
+                )
+                labels = torch.load(
+                    join(self.download_directory, self.dataset_name, "labels.pt")
+                )
+
+                self.dl_builder = DataLoaderBuilder(
+                    [
+                        FromArray(data, labels),
+                    ]
+                )
+            elif self.dataset_metadata["data_format"] == "numpy_array":
+                data = np.load(
+                    join(self.download_directory, self.dataset_name, "data.npy")
+                )
+                labels = np.load(
+                    join(self.download_directory, self.dataset_name, "labels.npy")
+                )
+                self.dl_builder = DataLoaderBuilder(
+                    [
+                        FromArray(data, labels),
+                    ]
+                )
             else:
-                raise ValueError(f"Data format {self.dataset_metadata['data_format']}" +
-                                 "is not yet supported.")
+                raise ValueError(
+                    f"Data format {self.dataset_metadata['data_format']}"
+                    + "is not yet supported."
+                )
         else:
-            raise ValueError(f"Dataset type {self.dataset_metadata['data_type']} is not yet supported.")
-
-    def _download_dataset(self,
-                          path_to_credentials: Union[None, str] = None,
-                          use_public_access: bool = True, ) -> None:
+            raise ValueError(
+                f"Dataset type {self.dataset_metadata['data_type']} is not yet supported."
+            )
+
+    def _download_dataset(
+        self,
+        path_to_credentials: Union[None, str] = None,
+        use_public_access: bool = True,
+    ) -> None:
         """Only download if the download directory does not exist already
         and if download directory contains at least three files (metadata,
         data, labels).
-        
+
         Args:
             path_to_credentials (str):
                 Path to the credentials file.
             use_public_access (bool):
                 Whether to use public access. If you want
                 to use the Google Cloud Storage API, you must set this to True.
-                
+
         Returns:
             None
         """
-        if (not os.path.isdir(join(self.download_directory, self.dataset_name))
-                or len(os.listdir(join(self.download_directory,
-                                       self.dataset_name))) < 3):
+        if (
+            not os.path.isdir(join(self.download_directory, self.dataset_name))
+            or len(os.listdir(join(self.download_directory, self.dataset_name))) < 3
+        ):
             # Delete the download directory if it exists but does not contain
             # the wanted number of files
-            if (os.path.isdir(join(self.download_directory, self.dataset_name))
-                    and
-                    len(os.listdir(join(self.download_directory,
-                                        self.dataset_name))) < 3):  # type: ignore
-                print("Deleting the download directory because it does " +
-                      "not contain the dataset")
+            if (
+                os.path.isdir(join(self.download_directory, self.dataset_name))
+                and len(os.listdir(join(self.download_directory, self.dataset_name)))
+                < 3
+            ):  # type: ignore
+                print(
+                    "Deleting the download directory because it does "
+                    + "not contain the dataset"
+                )
                 shutil.rmtree(self.download_directory, ignore_errors=True)
 
-            print("Downloading dataset {} to {}" \
-                  .format(self.dataset_name, self.download_directory))
-            dataset_cloud = DatasetCloud(self.dataset_name,
-                                         download_directory=self.download_directory,
-                                         path_to_credentials=path_to_credentials,
-                                         use_public_access=use_public_access,
-                                         )
+            print(
+                "Downloading dataset {} to {}".format(
+                    self.dataset_name, self.download_directory
+                )
+            )
+            dataset_cloud = DatasetCloud(
+                self.dataset_name,
+                download_directory=self.download_directory,
+                path_to_credentials=path_to_credentials,
+                use_public_access=use_public_access,
+            )
             dataset_cloud.download()
             del dataset_cloud
         else:
             print("Dataset '%s' already downloaded" % self.dataset_name)
 
     def get_metadata(self) -> Dict[str, Any]:
-        """ Returns the metadata of the dataset.
-        
+        """Returns the metadata of the dataset.
+
         Returns:
             Dict[str, Any]:
                 The metadata of the dataset.
         """
         return self.dataset_metadata
 
-    def build(self, tuple_of_kwargs: List[Dict[str, Any]]) \
-            -> Tuple[DataLoader, DataLoader, DataLoader]:
+    def build(
+        self, tuple_of_kwargs: List[Dict[str, Any]]
+    ) -> Tuple[DataLoader, DataLoader, DataLoader]:
         """Builds the dataloaders for the dataset.
-        
+
         Args:
             **tuple_of_kwargs: Arguments for the dataloader builder.
-            
+
         Returns:
             Tuple[DataLoader, DataLoader, DataLoader]:
                 The dataloaders for the dataset (train, validation, test).
diff --git a/gdeep/data/datasets/dataset_cloud.py b/gdeep/data/datasets/dataset_cloud.py
index e54f9d6d..ed87ca19 100644
--- a/gdeep/data/datasets/dataset_cloud.py
+++ b/gdeep/data/datasets/dataset_cloud.py
@@ -1,4 +1,3 @@
-
 import os
 from os import remove
 from os.path import join, exists
@@ -10,7 +9,8 @@
 from ._data_cloud import _DataCloud  # type: ignore
 from gdeep.utility.constants import DEFAULT_DOWNLOAD_DIR, DATASET_BUCKET_NAME
 
-class DatasetCloud():
+
+class DatasetCloud:
     """DatasetCloud class to handle the download and upload
     of datasets to the DataCloud.
     If the download_directory does not exist, it will be created and
@@ -37,25 +37,26 @@ class DatasetCloud():
             provided. Defaults to None.
         make_public (bool, optional):
             If True, the dataset will be made public
-            
+
         Raises:
             ValueError:
                 Dataset does not exits in cloud.
-        
+
         Returns:
             None
-        """
-    def __init__(self,
-             dataset_name: str,
-             bucket_name: str = DATASET_BUCKET_NAME,
-             download_directory: Union[None, str] = None,
-             use_public_access: bool = True,
-             path_to_credentials: Union[None, str] = None,
-             make_public: bool = True,
-             ) -> None:
+    """
+
+    def __init__(
+        self,
+        dataset_name: str,
+        bucket_name: str = DATASET_BUCKET_NAME,
+        download_directory: Union[None, str] = None,
+        use_public_access: bool = True,
+        path_to_credentials: Union[None, str] = None,
+        make_public: bool = True,
+    ) -> None:
         # Non-public datasets start with "private_"
-        if make_public or use_public_access or \
-            dataset_name.startswith("private_"):
+        if make_public or use_public_access or dataset_name.startswith("private_"):
             self.name = dataset_name
         else:
             self.name = "private_" + dataset_name
@@ -67,17 +68,17 @@ def __init__(self,
             self.download_directory = DEFAULT_DOWNLOAD_DIR
         else:
             self.download_directory = download_directory
-        
+
         self._data_cloud = _DataCloud(
             bucket_name=bucket_name,
-            download_directory = self.download_directory,
+            download_directory=self.download_directory,
             use_public_access=use_public_access,
-            path_to_credentials = path_to_credentials)
+            path_to_credentials=path_to_credentials,
+        )
         if use_public_access:
-            self.public_url = ("https://storage.googleapis.com/"
-                               + bucket_name + "/")
+            self.public_url = "https://storage.googleapis.com/" + bucket_name + "/"
         self.make_public = make_public
-        
+
     def __del__(self) -> None:
         """This function deletes the metadata file if it exists.
 
@@ -87,7 +88,7 @@ def __del__(self) -> None:
         if self.path_metadata != None:
             remove(self.path_metadata)  # type: ignore
         return None
-            
+
     def download(self) -> None:
         """Download a dataset from the DataCloud. If the dataset does not
         exist in the cloud, an exception will be raised. If the dataset
@@ -110,28 +111,32 @@ def _download_using_api(self) -> None:
         If the dataset does not exist in the bucket, an exception will
         be raised. If the dataset exists locally in the download directory,
         the dataset will not be downloaded again.
-        
+
         Raises:
             ValueError:
                 Dataset does not exits in cloud.
-        
+
         Returns:
             None
         """
         self._check_public_access()
         # List of existing datasets in the cloud.
-        existing_datasets: Set[str] = set([blob.name.split('/')[0]  # type: ignore
-                                 for blob in  # type: ignore
-                                 self._data_cloud.bucket.list_blobs()  # type: ignore
-            if blob.name != "giotto-deep-big.png"])  # type: ignore
+        existing_datasets: Set[str] = set(
+            [
+                blob.name.split("/")[0]  # type: ignore
+                for blob in self._data_cloud.bucket.list_blobs()  # type: ignore  # type: ignore
+                if blob.name != "giotto-deep-big.png"
+            ]
+        )  # type: ignore
         if self.name not in existing_datasets:
-            raise ValueError("Dataset {} does not exist in the cloud."\
-                .format(self.name) +
-                             "Available datasets are: {existing_datasets}.")
+            raise ValueError(
+                "Dataset {} does not exist in the cloud.".format(self.name)
+                + "Available datasets are: {existing_datasets}."
+            )
         if not self._does_dataset_exist_locally():
             self._create_dataset_folder()
-        self._data_cloud.download_folder(self.name + '/')
-    
+        self._data_cloud.download_folder(self.name + "/")
+
     def _does_dataset_exist_locally(self) -> bool:
         """Check if the dataset exists locally.
 
@@ -139,58 +144,63 @@ def _does_dataset_exist_locally(self) -> bool:
             bool: True if the dataset exists locally, False otherwise.
         """
         return exists(join(self.download_directory, self.name))
-    
+
     def _create_dataset_folder(self) -> None:
         """Creates a folder with the dataset name in the download directory.
-        
+
         Returns:
             None
         """
         if not exists(join(self.download_directory, self.name)):
             os.makedirs(join(self.download_directory, self.name))
-    
+
     def _download_using_url(self) -> None:
-        """ Download the dataset using the public url.
+        """Download the dataset using the public url.
         If the dataset does not exist in the bucket, an exception will
         be raised. If the dataset exists locally in the download directory,
         the dataset will not be downloaded again.
-        
+
         Raises:
             ValueError:
                 Dataset does not exits in cloud.
-        
+
         Returns:
             None
         """
         # List of existing datasets in the cloud.
         existing_datasets = self.get_existing_datasets()
-        
+
         # Check if requested dataset exists in the cloud.
-        assert self.name in existing_datasets,\
-            ("Dataset {} does not exist in the cloud.".format(self.name) +
-             "Available datasets are: {}.".format(existing_datasets))
-        
+        assert (
+            self.name in existing_datasets
+        ), "Dataset {} does not exist in the cloud.".format(
+            self.name
+        ) + "Available datasets are: {}.".format(
+            existing_datasets
+        )
+
         # If the dataset does not exist locally, create the dataset folder.
         if not self._does_dataset_exist_locally():
             self._create_dataset_folder()
-                    
-        # Download the dataset (metadata.json, data.pt, labels.pt) 
+
+        # Download the dataset (metadata.json, data.pt, labels.pt)
         # by using the public URL.
         self._data_cloud.download_file(self.name + "/metadata.json")
         # load the metadata.json file to get the filetype
-        with open(join(self.download_directory,  # type: ignore
-                       self.name, 'metadata.json')) as f:
+        with open(
+            join(self.download_directory, self.name, "metadata.json")  # type: ignore
+        ) as f:
             metadata = json.load(f)
         # filetype: Literal['pt', 'npy']
-        if metadata['data_format'] == "pytorch_tensor":
+        if metadata["data_format"] == "pytorch_tensor":
             filetype = "pt"
-        elif metadata['data_format'] == "numpy_array":
+        elif metadata["data_format"] == "numpy_array":
             filetype = "npy"
         else:
             raise ValueError(f"Unknown data format: {metadata['data_format']}")
         self._data_cloud.download_file(self.name + "/data." + filetype)
         self._data_cloud.download_file(self.name + "/labels." + filetype)
-        
+
     def get_existing_datasets(self) -> List[str]:
         """Returns a list of datasets in the cloud.
 
@@ -201,83 +211,87 @@ def get_existing_datasets(self) -> List[str]:
         if self.use_public_access:
             datasets_local = "tmp_datasets.json"
             # Download the dataset list json file using the public URL.
-            wget.download(self.public_url + 'datasets.json', datasets_local)  # type: ignore
+            wget.download(self.public_url + "datasets.json", datasets_local)  # type: ignore
             datasets = json.load(open(datasets_local))
-            
+
             # Remove duplicates. This has to be fixed in the future.
             datasets = list(set(datasets))
-            
+
             # Remove the temporary file.
             remove(datasets_local)
-            
+
             return datasets
         else:
-            existing_datasets = [blob_name.split('/')[0] 
-                    for blob_name in self._data_cloud.list_blobs()
-                    if blob_name != "giotto-deep-big.png" and
-                    blob_name != "datasets.json"]
+            existing_datasets = [
+                blob_name.split("/")[0]
+                for blob_name in self._data_cloud.list_blobs()
+                if blob_name != "giotto-deep-big.png" and blob_name != "datasets.json"
+            ]
             # Remove duplicates.
             existing_datasets = list(set(existing_datasets))
-            
+
             # Remove dataset that are not public, i.e. start with "private_".
-            existing_datasets = [dataset for dataset in existing_datasets
-                                    if not dataset.startswith("private_")]
-            
+            existing_datasets = [
+                dataset
+                for dataset in existing_datasets
+                if not dataset.startswith("private_")
+            ]
+
             return existing_datasets
 
     def _update_dataset_list(self) -> None:
         """Updates the dataset list in the datasets.json file.
-        
+
         Returns:
             None
         """
         self._check_public_access()
-        
+
         # List of existing datasets in the cloud.
         existing_datasets = self.get_existing_datasets()
-        
+
         # Save existing datasets to a json file.
-        json_file = 'tmp_datasets.json'
-        json.dump(existing_datasets, open(json_file, 'w'))
-        
+        json_file = "tmp_datasets.json"
+        json.dump(existing_datasets, open(json_file, "w"))
+
         # Upload the json file to the cloud.
-        self._data_cloud.upload_file(json_file,
-                                     'datasets.json',
-                                     make_public=True,
-                                     overwrite=True,
-                                     )
-        
+        self._data_cloud.upload_file(
+            json_file,
+            "datasets.json",
+            make_public=True,
+            overwrite=True,
+        )
+
         # Remove the temporary file.
         remove(json_file)
-        
-        
-        
 
     @staticmethod
     def _get_filetype(path: str) -> str:
         """Returns the file extension from a given path.
-    
+
         Args:
             path:
                 A string path.
-        
+
         Returns:
-            str: 
+            str:
                 The file extension.
-        
+
         Raises:
             None.
         """
-        return path.split('.')[-1]
-    
+        return path.split(".")[-1]
+
     def _check_public_access(self) -> None:
-        """Check if use_public_access is set to False.
-        """
-        assert self.use_public_access is False,\
-            "Only download functionality is supported for public access."
-    
-    def _upload_data(self,
-                path: str,) -> None:
+        """Check if use_public_access is set to False."""
+        assert (
+            self.use_public_access is False
+        ), "Only download functionality is supported for public access."
+
+    def _upload_data(
+        self,
+        path: str,
+    ) -> None:
         """Uploads the data file to a Cloud Storage bucket.
 
         Args:
@@ -287,33 +301,35 @@ def _upload_data(self,
         Raises:
         ValueError:
             If the file type is not supported.
-        
+
         Returns:
             None
         """
         self._check_public_access()
-        
+
         filetype = DatasetCloud._get_filetype(path)
-        
+
         # Check if the file type is supported
-        if filetype in ['pt', 'npy']:
-            self._data_cloud.upload_file(path,
-                                         (self.metadata['name'] + # type: ignore
-                                            '/data.' + filetype),
-                                         make_public=self.make_public,
-                                         overwrite=False,
-                                         )
+        if filetype in ["pt", "npy"]:
+            self._data_cloud.upload_file(
+                path,
+                (self.metadata["name"] + "/data." + filetype),  # type: ignore
+                make_public=self.make_public,
+                overwrite=False,
+            )
         else:
             raise ValueError("File type {} is not supported.".format(filetype))
-    
-    def _upload_label(self,
-                 path: str,) -> None:
+
+    def _upload_label(
+        self,
+        path: str,
+    ) -> None:
         """Uploads a set of labels to a remote dataset.
 
         Args:
             path:
                 the path to the labels file.
-        
+
         Raises:
             ValueError:
                 If the file type is not supported.
@@ -323,55 +339,55 @@ def _upload_label(self,
 
         """
         self._check_public_access()
-        
+
         filetype = DatasetCloud._get_filetype(path)
-        
+
         # Check if the file type is supported
-        if filetype in ['pt', 'npy']:
-            self._data_cloud.upload_file(path,
-                                         (self.metadata['name'] + # type: ignore
-                                            '/labels.' + filetype),
-                                         make_public=self.make_public,
-                                         overwrite=False,
-                                         )
+        if filetype in ["pt", "npy"]:
+            self._data_cloud.upload_file(
+                path,
+                (self.metadata["name"] + "/labels." + filetype),  # type: ignore
+                make_public=self.make_public,
+                overwrite=False,
+            )
         else:
             raise ValueError("File type {} is not supported.".format(filetype))
 
-    
-    def _upload_metadata(self,
-                        path: Union[str, None]=None) -> None:
+    def _upload_metadata(self, path: Union[str, None] = None) -> None:
         """Uploads the metadata dictionary to the location specified in the
         metadata. The metadata dictionary is generated using create_metadata.
-        
+
         Args:
             path (str):
                 The path to the data cloud folder. If none, path will
                 be set to the default path.
-            
+
         Raises:
             Exception:
                 If no metadata exists, an exception will be raised.
-                
+
         Returns:
             None
         """
         self._check_public_access()
         self._data_cloud.upload_file(
-            path, str(self.metadata['name']) + '/'  # type: ignore
-            + 'metadata.json',
-            make_public=self.make_public)
-        
-    def _add_metadata(self,
-                     size_dataset: int,
-                     input_size: Tuple[int, ...],
-                     num_labels: Union[None, int] = None,
-                     data_type: str = "tabular",
-                     task_type: str = "classification",
-                     name: Union[None, str]=None,
-                     data_format: Union[None, str]=None,
-                     comment: Union[None, str]=None,
-                     ) -> None:
-        """This function accepts various metadata for the dataset and stores it 
+            path,
+            str(self.metadata["name"]) + "/" + "metadata.json",  # type: ignore
+            make_public=self.make_public,
+        )
+
+    def _add_metadata(
+        self,
+        size_dataset: int,
+        input_size: Tuple[int, ...],
+        num_labels: Union[None, int] = None,
+        data_type: str = "tabular",
+        task_type: str = "classification",
+        name: Union[None, str] = None,
+        data_format: Union[None, str] = None,
+        comment: Union[None, str] = None,
+    ) -> None:
+        """This function accepts various metadata for the dataset and stores it
         in a temporary JSON file.
 
         Args:
@@ -379,7 +395,7 @@ def _add_metadata(self,
                 The size of the dataset (in terms of the number
                 of samples).
             input_size (Tuple[int, ...]):
-                The size of each sample in the 
+                The size of each sample in the
                 dataset.
             num_labels (Union[None, int]):
                 The number of classes in the dataset.
@@ -403,23 +419,25 @@ def _add_metadata(self,
         if data_format is None:
             data_format = "pytorch_tensor"
         self.path_metadata = "tmp_metadata.json"  # type: ignore
-        self.metadata = {'name': name,
-                    'size': size_dataset,
-                    'input_size': input_size,
-                    'num_labels': num_labels,
-                    'task_type': task_type,
-                    'data_type': data_type,
-                    'data_format': data_format,
-                    'comment': comment
-                    }
+        self.metadata = {
+            "name": name,
+            "size": size_dataset,
+            "input_size": input_size,
+            "num_labels": num_labels,
+            "task_type": task_type,
+            "data_type": data_type,
+            "data_format": data_format,
+            "comment": comment,
+        }
         with open(self.path_metadata, "w") as f:  # type: ignore
             json.dump(self.metadata, f, sort_keys=True, indent=4)
-        
-    def _upload(self,
-                       path_data: str,
-                       path_label: str,
-                       path_metadata: Union[str, None] = None,
-                       ) -> None:
+
+    def _upload(
+        self,
+        path_data: str,
+        path_label: str,
+        path_metadata: Union[str, None] = None,
+    ) -> None:
         """Uploads a dataset to the cloud.
 
         Args:
@@ -429,23 +447,24 @@ def _upload(self,
 
         Raises:
             ValueError: If the dataset already exists in the cloud.
-        
+
         Returns:
-            None    
+            None
         """
         self._check_public_access()
-        
+
         # List of existing datasets in the cloud.
         existing_datasets = self.get_existing_datasets()
         if self.name in existing_datasets:
-            raise ValueError("Dataset {} already exists in the cloud."\
-                .format(self.name) +
-                "Available datasets are: {}.".format(existing_datasets))
+            raise ValueError(
+                "Dataset {} already exists in the cloud.".format(self.name)
+                + "Available datasets are: {}.".format(existing_datasets)
+            )
         if path_metadata is None:
             path_metadata = self.path_metadata
         self._upload_metadata(path_metadata)
         self._upload_data(path_data)
         self._upload_label(path_label)
-        
+
         # Update dataset list.
-        self._update_dataset_list()
\ No newline at end of file
+        self._update_dataset_list()
diff --git a/gdeep/data/datasets/dataset_for_image.py b/gdeep/data/datasets/dataset_for_image.py
index 176c3a6e..0b770101 100644
--- a/gdeep/data/datasets/dataset_for_image.py
+++ b/gdeep/data/datasets/dataset_for_image.py
@@ -19,15 +19,14 @@
 from tqdm import tqdm
 
 
-
 Tensor = torch.Tensor
-T = TypeVar('T')
+T = TypeVar("T")
 
 
 class ImageClassificationFromFiles(Dataset[Any]):
     """This class is useful to build a dataset
     directly from image files
-    
+
     Args:
         img_folder (string):
             The path to the folder where the training
@@ -49,19 +48,20 @@ class ImageClassificationFromFiles(Dataset[Any]):
     img_folder: str
     labels_file: str
     img_labels: pd.DataFrame
-    
-    def __init__(self, img_folder: str=".",
-                 labels_file:str="labels.csv",
-                 ) -> None:
+
+    def __init__(
+        self,
+        img_folder: str = ".",
+        labels_file: str = "labels.csv",
+    ) -> None:
 
         self.img_folder = img_folder
         self.img_labels = pd.read_csv(labels_file)
 
-
     def __len__(self) -> int:
         return len(self.img_labels)
 
-    def __getitem__(self, idx:int) -> Tuple[Any, Union[str, int]]:
+    def __getitem__(self, idx: int) -> Tuple[Any, Union[str, int]]:
         image = self._get_image(idx)
         image_out = deepcopy(image)
         label = self.img_labels.iloc[idx, 1]
@@ -79,4 +79,3 @@ def _get_image(self, idx: int) -> Any:
             warnings.warn(f"The image {img_path} canot be loaded. Skipping it.")
             return None
         return image
-
diff --git a/gdeep/data/datasets/dataset_form_array.py b/gdeep/data/datasets/dataset_form_array.py
index b3ee1992..ac811bd4 100644
--- a/gdeep/data/datasets/dataset_form_array.py
+++ b/gdeep/data/datasets/dataset_form_array.py
@@ -1,8 +1,6 @@
-
 import os
 import shutil
-from typing import Any, Callable, Dict, Optional, \
-    Optional, Tuple, TypeVar, Union
+from typing import Any, Callable, Dict, Optional, Optional, Tuple, TypeVar, Union
 
 import numpy as np
 import torch
@@ -26,9 +24,8 @@ class FromArray(Dataset[Tuple[Tensor, Tensor]]):
             with the data
 
     """
-    def __init__(self, x: Union[Tensor, Array],
-                 y: Union[Tensor, Array]
-                 ) -> None:
+
+    def __init__(self, x: Union[Tensor, Array], y: Union[Tensor, Array]) -> None:
         self.x = self._from_numpy(x)
         y = self._from_numpy(y)
         self.y = self._long_or_float(y)
@@ -48,13 +45,13 @@ def _long_or_float(y: Union[Tensor, Array]) -> Tensor:
         a long tensor of a float tensor"""
         if isinstance(y, torch.Tensor):
             return y
-        if y.dtype in [np.float16, np.float32, np.float64]:
+        if y.dtype in [np.float16, np.float32, np.float64]:  # type: ignore
             return torch.tensor(y).float()
         return torch.tensor(y).long()
 
     def __len__(self):
         return self.y.shape[0]
 
-    def __getitem__(self, idx:int) -> Tuple[Tensor, Tensor]:
+    def __getitem__(self, idx: int) -> Tuple[Tensor, Tensor]:
         x, y = (self.x[idx], self.y[idx])
         return x, y
diff --git a/gdeep/data/datasets/graph_datasets.py b/gdeep/data/datasets/graph_datasets.py
index d4c43a3c..9ed9ea4c 100644
--- a/gdeep/data/datasets/graph_datasets.py
+++ b/gdeep/data/datasets/graph_datasets.py
@@ -2,5 +2,3 @@
 
 
 from gdeep.utility.constants import ROOT_DIR, DEFAULT_GRAPH_DIR
-
-
diff --git a/gdeep/data/datasets/orbit.py b/gdeep/data/datasets/orbit.py
index 9034998d..9a40e6b2 100644
--- a/gdeep/data/datasets/orbit.py
+++ b/gdeep/data/datasets/orbit.py
@@ -2,40 +2,41 @@
 import warnings
 
 
-def generate_orbit(num_pts_per_orbit: int = 100,
-                   parameter: float = 1.0) -> np.ndarray:
+def generate_orbit(num_pts_per_orbit: int = 100, parameter: float = 1.0) -> np.ndarray:
     """Generate sequence of points of a dynamical system
     Non-parallel version
 
     Args:
-        num_pts_per_orbit (int, optional): 
+        num_pts_per_orbit (int, optional):
             Number of points to generate.
             Defaults to 100.
-        parameter (float, optional): 
+        parameter (float, optional):
             Parameter of the dynamical system.
             Defaults to 1.0.
 
     Returns:
-        np.ndarray: 
+        np.ndarray:
             Array of sampled points of the dynamical system.
     """
-    warnings.warn("This methode is deprecated. Use the parallel\
+    warnings.warn(
+        "This methode is deprecated. Use the parallel\
                    version create_data.parallel_orbit.generate_\
                        orbit_parallel",
-                  DeprecationWarning)
+        DeprecationWarning,
+    )
     try:
         assert parameter > 0
     except AssertionError:
-        print('Parameter must be greater than 0')
+        print("Parameter must be greater than 0")
     try:
         assert num_pts_per_orbit > 0
     except AssertionError:
-        print('num_pts_per_orbit must be greater than 0')
+        print("num_pts_per_orbit must be greater than 0")
 
     orbit = np.zeros([num_pts_per_orbit, 2])
     x_cur, y_cur = np.random.rand(), np.random.rand()  # current points
     for idx in range(num_pts_per_orbit):
-        x_cur = (x_cur + parameter * y_cur * (1. - y_cur)) % 1
-        y_cur = (y_cur + parameter * x_cur * (1. - x_cur)) % 1
+        x_cur = (x_cur + parameter * y_cur * (1.0 - y_cur)) % 1
+        y_cur = (y_cur + parameter * x_cur * (1.0 - x_cur)) % 1
         orbit[idx, :] = [x_cur, y_cur]
     return orbit
diff --git a/gdeep/data/datasets/parallel_orbit.py b/gdeep/data/datasets/parallel_orbit.py
index 4dd3864c..f25c5dd7 100644
--- a/gdeep/data/datasets/parallel_orbit.py
+++ b/gdeep/data/datasets/parallel_orbit.py
@@ -8,22 +8,25 @@
 from gtda.homology import WeakAlphaPersistence  # type: ignore
 from sympy import N, S
 
+
 class DataLoaderKwargs(object):
-    """Object to store keyword arguments for train, val, and test dataloaders
-    """
+    """Object to store keyword arguments for train, val, and test dataloaders"""
+
     def __init__(self, *, train_kwargs, val_kwargs, test_kwargs):
         self.train_kwargs = train_kwargs
         self.val_kwargs = val_kwargs
         self.test_kwargs = test_kwargs
-        
+
     def get_train_args(self):
         return self.train_kwargs
-    
+
     def get_val_args(self):
         return self.val_kwargs
-    
+
     def get_test_args(self):
         return self.test_kwargs
+
+
 class OrbitsGenerator(object):
     """Generate Orbit dataset consistent of orbits defined by the dynamical system
     x[n+1] = x[n] + r * y[n] * (1  - y[n]) % 1
@@ -33,39 +36,41 @@ class OrbitsGenerator(object):
     it given the orbit.
     By default r is chosen from (2.5, 3.5, 4.0, 4.1, 4.3).
     Args:
-        parameters (Tuple[float]): 
+        parameters (Tuple[float]):
             Hyperparameter of the dynamical systems.
-        num_orbits_per_class (int): 
+        num_orbits_per_class (int):
             number of orbits per class.
-        num_pts_per_orbit (int): 
+        num_pts_per_orbit (int):
             number of points per orbit.
-        homology_dimensions (Sequence[int]): 
+        homology_dimensions (Sequence[int]):
             homology dimension of the persistence
             diagrams.
-        validation_percentage (float, optional): 
+        validation_percentage (float, optional):
             Percentage of the validation dataset.
             Defaults to 0.0.
-        test_percentage (float, optional): 
+        test_percentage (float, optional):
             Percentage of the test dataset. Defaults to 0.0.
-        dynamical_system (str, optional): 
+        dynamical_system (str, optional):
             either use persistence paths
             convention ´pp_convention´ or the classical convention
             ´classical_convention´. Defaults to '´classical_convention´'.
-        n_jobs (int, optional): 
+        n_jobs (int, optional):
             number of cpus to run the computation on. Defaults to 1.
     """
-    def __init__(self,
-                 parameters: Sequence[float] = (2.5, 3.5, 4.0, 4.1, 4.3),
-                 num_orbits_per_class: int = 1_000,
-                 num_pts_per_orbit: int = 1_000,
-                 homology_dimensions: Sequence[int] = (0, 1),
-                 validation_percentage: float = 0.0,
-                 test_percentage: float = 0.0,
-                 dynamical_system: str = 'classical_convention',
-                 n_jobs: int = 1,
-                 dtype: str = 'float32',
-                 arbitrary_precision=False,
-                 ) -> None:
+
+    def __init__(
+        self,
+        parameters: Sequence[float] = (2.5, 3.5, 4.0, 4.1, 4.3),
+        num_orbits_per_class: int = 1_000,
+        num_pts_per_orbit: int = 1_000,
+        homology_dimensions: Sequence[int] = (0, 1),
+        validation_percentage: float = 0.0,
+        test_percentage: float = 0.0,
+        dynamical_system: str = "classical_convention",
+        n_jobs: int = 1,
+        dtype: str = "float32",
+        arbitrary_precision=False,
+    ) -> None:
 
         # Initialize member variables.
         self._parameters = parameters
@@ -76,76 +81,93 @@ def __init__(self,
         self._num_homology_dimensions = len(self._homology_dimensions)
         self._n_jobs = n_jobs
         # Assert that validation and testing parameters are valid
-        assert (test_percentage >= 0.0 and validation_percentage >= 0.0
-                and test_percentage + validation_percentage < 1.0)
+        assert (
+            test_percentage >= 0.0
+            and validation_percentage >= 0.0
+            and test_percentage + validation_percentage < 1.0
+        )
         self._validation_percentage = validation_percentage
         self._test_percentage = test_percentage
-        
+
         # Assert that convention for the dynamical system is valid
-        assert dynamical_system in ('classical_convention', 'pp_convention'),\
-            f'{dynamical_system} is not supported'
+        assert dynamical_system in (
+            "classical_convention",
+            "pp_convention",
+        ), f"{dynamical_system} is not supported"
         self.dynamical_system = dynamical_system
-        
+
         # Initialize orbits array and persistence diagrams array with None.
         self._orbits = None
         self._labels = None
         self._persistence_diagrams = None
-        
+
         # Initialize the train, val, and test indices with None.
         self._train_idcs = None
         self._val_idcs = None
         self._test_idcs = None
 
-        assert dtype in ('float32', 'float64', 'float128'), f"Type {dtype} is not supported."
+        assert dtype in (
+            "float32",
+            "float64",
+            "float128",
+        ), f"Type {dtype} is not supported."
         self._dtype = dtype
 
         self.arbitrary_precision = arbitrary_precision
 
     def orbits_from_array(self, orbits):
-        assert (orbits.shape[0] == self._num_orbits_per_class * self._num_classes and
-                orbits.shape[1] == self._num_pts_per_orbit and
-                orbits.shape[2] == 2), "Array does not have the right shape."
+        assert (
+            orbits.shape[0] == self._num_orbits_per_class * self._num_classes
+            and orbits.shape[1] == self._num_pts_per_orbit
+            and orbits.shape[2] == 2
+        ), "Array does not have the right shape."
 
         self._orbits = orbits
 
-        y = np.array([self._num_orbits_per_class * [c]
-                      for c in range(self._num_classes)])
+        y = np.array(
+            [self._num_orbits_per_class * [c] for c in range(self._num_classes)]
+        )
 
         self._labels = y.reshape(-1)
 
     def _generate_orbits(self) -> None:
-        """Generate the orbits for the dynamical system.
-        """
-        
+        """Generate the orbits for the dynamical system."""
+
         # If orbits are already computed do nothing
         if self._orbits is not None:
             return
         else:
             # Initialize the orbits array with zeros.
-            x = np.zeros((
-                self._num_classes,  # type: ignore
-                self._num_orbits_per_class,
-                self._num_pts_per_orbit,
-                2
-            ))
+            x = np.zeros(
+                (
+                    self._num_classes,  # type: ignore
+                    self._num_orbits_per_class,
+                    self._num_pts_per_orbit,
+                    2,
+                )
+            )
             # Initialize the labels array with the hyperparameter indices.
-            y = np.array([self._num_orbits_per_class * [c]
-                          for c in range(self._num_classes)])
-            
-            self._labels = y.reshape(-1).astype('int64')  # type: ignore
+            y = np.array(
+                [self._num_orbits_per_class * [c] for c in range(self._num_classes)]
+            )
+
+            self._labels = y.reshape(-1).astype("int64")  # type: ignore
             # generate dataset
             for class_idx, p in enumerate(self._parameters):  # type: ignore
                 x[class_idx, :, 0, :] = np.random.rand(self._num_orbits_per_class, 2)  # type: ignore
 
                 if self.arbitrary_precision:
-                    assert self.dynamical_system == 'classical_convention', "Only classical_convention implemented yet"
+                    assert (
+                        self.dynamical_system == "classical_convention"
+                    ), "Only classical_convention implemented yet"
                     for orbit in range(self._num_orbits_per_class):
                         print(orbit)
                         x[class_idx, orbit, :, :] = self._orbit_high_precision(
-                                                        x_init=x[class_idx, orbit, 0, :],
-                                                        rho=p,
-                                                        num_points=self._num_pts_per_orbit,
-                                                        precision=600)
+                            x_init=x[class_idx, orbit, 0, :],
+                            rho=p,
+                            num_points=self._num_pts_per_orbit,
+                            precision=600,
+                        )
 
                 else:
 
@@ -153,21 +175,27 @@ def _generate_orbits(self) -> None:
                         x_cur = x[class_idx, :, i - 1, 0]
                         y_cur = x[class_idx, :, i - 1, 1]
 
-                        if self.dynamical_system == 'pp_convention':
-                            x[class_idx, :, i, 0] = (x_cur + p * y_cur * (1. - y_cur)) % 1
-                            x[class_idx, :, i, 1] = (y_cur + p * x_cur * (1. - x_cur)) % 1
+                        if self.dynamical_system == "pp_convention":
+                            x[class_idx, :, i, 0] = (
+                                x_cur + p * y_cur * (1.0 - y_cur)
+                            ) % 1
+                            x[class_idx, :, i, 1] = (
+                                y_cur + p * x_cur * (1.0 - x_cur)
+                            ) % 1
                         else:
-                            x[class_idx, :, i, 0] = (x_cur + p * y_cur * (1. - y_cur)) % 1
+                            x[class_idx, :, i, 0] = (
+                                x_cur + p * y_cur * (1.0 - y_cur)
+                            ) % 1
                             x_next = x[class_idx, :, i, 0]
-                            x[class_idx, :, i, 1] = (y_cur + p * x_next * (1. - x_next)) % 1
+                            x[class_idx, :, i, 1] = (
+                                y_cur + p * x_next * (1.0 - x_next)
+                            ) % 1
 
-
-            self._orbits = x.reshape((-1, self._num_pts_per_orbit, 2))  #type: ignore
+            self._orbits = x.reshape((-1, self._num_pts_per_orbit, 2))  # type: ignore
 
     def _orbit_high_precision(self, x_init, rho, num_points=1_000, precision=600):
         x_precise = np.zeros((1_000, 2))
 
-
         x0 = S(x_init[0])
         y0 = S(x_init[1])
 
@@ -181,143 +209,153 @@ def _orbit_high_precision(self, x_init, rho, num_points=1_000, precision=600):
         return x_precise
 
     def _compute_persistence_diagrams(self) -> None:
-        """ Computes the weak alpha persistence of the orbit data clouds.
-            The result is stored in member variable ´persistence_diagrams´.
-            It has the shape
-            [num_classes, num_orbits, num_persistence_points, 3].
-            In the last dimension the first two values are the coordinates of
-            the points in the persistence diagrams and the third is the
-            homology dimension.
+        """Computes the weak alpha persistence of the orbit data clouds.
+        The result is stored in member variable ´persistence_diagrams´.
+        It has the shape
+        [num_classes, num_orbits, num_persistence_points, 3].
+        In the last dimension the first two values are the coordinates of
+        the points in the persistence diagrams and the third is the
+        homology dimension.
         """
         if self._orbits is None:
             self._generate_orbits()
-        
+
         # compute weak alpha persistence
         wap = WeakAlphaPersistence(
-                            homology_dimensions=self._homology_dimensions,
-                            n_jobs=self._n_jobs
-                            )
+            homology_dimensions=self._homology_dimensions, n_jobs=self._n_jobs
+        )
         # c: class, o: orbit, p: point, d: dimension
-        #orbits_stack = rearrange(self._orbits, 'c o p d -> (c o) p d')  # stack classes
+        # orbits_stack = rearrange(self._orbits, 'c o p d -> (c o) p d')  # stack classes
         persistence_diagrams_categorical = wap.fit_transform(self._orbits)
         # shape: (num_classes * n_samples, n_features, 3)
         # Convert persistence diagram to one-hot homological dimension encoding
         self._persistence_diagrams = self._persistence_diagrams_to_one_hot(
-                                        persistence_diagrams_categorical
-                                        ).astype(self._dtype)
+            persistence_diagrams_categorical
+        ).astype(self._dtype)
 
     def _persistence_diagrams_to_one_hot(self, persistence_diagrams):
-        """ Convert homology dimension to one-hot encoding
+        """Convert homology dimension to one-hot encoding
         Args:
-            persistence_diagrams ([np.array]): 
+            persistence_diagrams ([np.array]):
                 persistence diagram with categorical
                 homology dimension.
         Returns:
-            [np.array]: 
+            [np.array]:
                 persistent diagram with one-hot encoded homology dimension.
         """
         if self._num_homology_dimensions > 1:
             return np.concatenate(
                 (
                     persistence_diagrams[:, :, :2],  # point coordinates
-                    (np.eye(self._num_homology_dimensions)  # type: ignore
-                    [persistence_diagrams[:, :, -1].astype(np.int32)]),
+                    (
+                        np.eye(self._num_homology_dimensions)[  # type: ignore
+                            persistence_diagrams[:, :, -1].astype(np.int32)
+                        ]
+                    ),
                 ),
-                axis=-1)
+                axis=-1,
+            )
         else:
             return persistence_diagrams[:, :, :2]
-    
+
     def get_orbits(self) -> Union[None, np.ndarray]:
         """Returns the orbits as an ndarrays of shape
         (num_classes * num_orbits_per_class, num_pts_per_orbit, 2)
         Returns:
-            np.ndarray: 
+            np.ndarray:
                 Orbits
         """
         if self._orbits is None:
             self._generate_orbits()
         return self._orbits
-    
+
     def get_persistence_diagrams(self) -> Union[None, np.ndarray]:
         """Returns the orbits as an ndarrays of shape
         (num_classes * num_orbits_per_class, num_topological_features, 3)
         Returns:
-            np.ndarray: 
+            np.ndarray:
                 Persistence diagrams
         """
         if self._persistence_diagrams is None:
             self._compute_persistence_diagrams()
-            
+
         return self._persistence_diagrams
-    
+
     def _split_data_idcs(self) -> None:
-        """Split the data indices into training, validation and testing data.
-        """
-        idcs = np.arange(self._num_classes
-                            * self._num_orbits_per_class)
+        """Split the data indices into training, validation and testing data."""
+        idcs = np.arange(self._num_classes * self._num_orbits_per_class)
 
         if self._test_percentage > 0.0:
-            rest_idcs, self._test_idcs = train_test_split(idcs,
-                                                        test_size=self._test_percentage)
+            rest_idcs, self._test_idcs = train_test_split(
+                idcs, test_size=self._test_percentage
+            )
         else:
             rest_idcs, self._test_idcs = idcs, []  # type: ignore
-            
+
         if self._validation_percentage > 0.0:
-            self._train_idcs, self._val_idcs = train_test_split(rest_idcs,
-                                                    test_size =(
-                                                    self._validation_percentage
-                                                    / (1.0 - self._test_percentage))
-                                                    )
+            self._train_idcs, self._val_idcs = train_test_split(
+                rest_idcs,
+                test_size=(self._validation_percentage / (1.0 - self._test_percentage)),
+            )
         else:
             self._train_idcs, self._val_idcs = rest_idcs, []  # type: ignore
-    
-    def _get_data_loaders(self, list_of_arrays: List[np.ndarray],
-                          dataloaders_kwargs: DataLoaderKwargs
-                          ) -> Tuple[DataLoader, DataLoader, DataLoader]:
+
+    def _get_data_loaders(
+        self, list_of_arrays: List[np.ndarray], dataloaders_kwargs: DataLoaderKwargs
+    ) -> Tuple[DataLoader, DataLoader, DataLoader]:
         """Generates a DataLoader for the given list of arrays.
         Args:
-            list_of_arrays ([List[np.ndarray]]): 
+            list_of_arrays ([List[np.ndarray]]):
                 List of arrays to load.
 
         Returns:
-            Tuple[DataLoader, DataLoader, DataLoader]: 
+            Tuple[DataLoader, DataLoader, DataLoader]:
                 train, val, test data loaders.
         """
-        assert ((self._train_idcs is not None) and 
-                (self._val_idcs is not None) and
-                (self._test_idcs is not None)),\
-                "Train, validation, and test data must be initialized"
-        
-        dl_train = DataLoader(TensorDataset(*(torch.tensor(a[self._train_idcs])
-                                              for a in list_of_arrays)),
-                                **dataloaders_kwargs.get_train_args())
-        dl_val = DataLoader(TensorDataset(*(torch.tensor(a[self._val_idcs])
-                                            for a in list_of_arrays)),
-                                **dataloaders_kwargs.get_val_args())
-        dl_test = DataLoader(TensorDataset(*(torch.tensor(a[self._train_idcs])
-                                             for a in list_of_arrays)),
-                                **dataloaders_kwargs.get_test_args())
+        assert (
+            (self._train_idcs is not None)
+            and (self._val_idcs is not None)
+            and (self._test_idcs is not None)
+        ), "Train, validation, and test data must be initialized"
+
+        dl_train = DataLoader(
+            TensorDataset(*(torch.tensor(a[self._train_idcs]) for a in list_of_arrays)),
+            **dataloaders_kwargs.get_train_args(),
+        )
+        dl_val = DataLoader(
+            TensorDataset(*(torch.tensor(a[self._val_idcs]) for a in list_of_arrays)),
+            **dataloaders_kwargs.get_val_args(),
+        )
+        dl_test = DataLoader(
+            TensorDataset(*(torch.tensor(a[self._train_idcs]) for a in list_of_arrays)),
+            **dataloaders_kwargs.get_test_args(),
+        )
         return dl_train, dl_val, dl_test
-    
-    def get_dataloader_orbits(self,
-                              dataloaders_kwargs: DataLoaderKwargs
-                              )-> Tuple[DataLoader, DataLoader, DataLoader]:
-        """Generates a Dataloader from the orbits dataset 
+
+    def get_dataloader_orbits(
+        self, dataloaders_kwargs: DataLoaderKwargs
+    ) -> Tuple[DataLoader, DataLoader, DataLoader]:
+        """Generates a Dataloader from the orbits dataset
         Returns:
-            DataLoader: 
+            DataLoader:
                 Dataloader of orbits
         """
         if self._orbits is None:
             self._generate_orbits()
         if self._train_idcs is None:
             self._split_data_idcs()
-        return self._get_data_loaders([self._orbits.astype(self._dtype),  # type: ignore
-                                            self._labels], # type: ignore
-                                        dataloaders_kwargs)
-    
-    def get_dataloader_persistence_diagrams(self, dataloaders_kwargs: DataLoaderKwargs
-                              )-> Tuple[DataLoader, DataLoader, DataLoader]:
-        """Generates a Dataloader from the persistence diagrams dataset 
+        return self._get_data_loaders(
+            [
+                self._orbits.astype(self._dtype),  # type: ignore
+                self._labels,
+            ],  # type: ignore
+            dataloaders_kwargs,
+        )
+
+    def get_dataloader_persistence_diagrams(
+        self, dataloaders_kwargs: DataLoaderKwargs
+    ) -> Tuple[DataLoader, DataLoader, DataLoader]:
+        """Generates a Dataloader from the persistence diagrams dataset
         Returns:
             Tuple[DataLoader, DataLoader, DataLoader]:
                 Dataloaders of persistence diagrams
@@ -327,24 +365,32 @@ def get_dataloader_persistence_diagrams(self, dataloaders_kwargs: DataLoaderKwar
         if self._train_idcs is None:
             self._split_data_idcs()
 
-        return self._get_data_loaders([self._persistence_diagrams, self._labels], # type: ignore
-                                        dataloaders_kwargs)
-    
-    def get_dataloader_combined(self, dataloaders_kwargs: DataLoaderKwargs,
-                              )-> Tuple[DataLoader, DataLoader, DataLoader]:
+        return self._get_data_loaders(
+            [self._persistence_diagrams, self._labels],  # type: ignore
+            dataloaders_kwargs,
+        )
+
+    def get_dataloader_combined(
+        self,
+        dataloaders_kwargs: DataLoaderKwargs,
+    ) -> Tuple[DataLoader, DataLoader, DataLoader]:
         """Generates a Dataloader from the orbits dataset and the persistence diagrams
         Returns:
-            DataLoader: 
+            DataLoader:
                 Dataloader of orbits and persistence diagrams
         """
         if self._persistence_diagrams is None:
             self._compute_persistence_diagrams()
         if self._train_idcs is None:
             self._split_data_idcs()
-        return self._get_data_loaders([self._orbits,  # type: ignore
-                                       self._persistence_diagrams,  # type: ignore
-                                       self._labels],  # type: ignore
-                                       dataloaders_kwargs)
+        return self._get_data_loaders(
+            [
+                self._orbits,  # type: ignore
+                self._persistence_diagrams,  # type: ignore
+                self._labels,
+            ],  # type: ignore
+            dataloaders_kwargs,
+        )
 
 
 def generate_orbit_parallel(
@@ -352,43 +398,38 @@ def generate_orbit_parallel(
     num_orbits,
     num_pts_per_orbit: int = 100,
     parameters: List[float] = [1.0],
-    ) -> np.ndarray:
+) -> np.ndarray:
     """Generate sequence of points of a dynamical system
     in a parallel manner.
 
     Args:
-        num_classes (int): 
+        num_classes (int):
             number of classes of dynamical systems.
-        num_orbits (int): 
+        num_orbits (int):
             number of orbits of dynamical system per class.
-        num_pts_per_orbit (int, optional): 
+        num_pts_per_orbit (int, optional):
             Number of points to generate.
             Defaults to 100.
-        parameter (List[float], optional): 
+        parameter (List[float], optional):
             List of parameters of the dynamical
             system.
             Defaults to [1.0].
 
     Returns:
-        np.ndarray: 
+        np.ndarray:
             Array of sampled points of the dynamical system.
     """
     try:
         for parameter in parameters:
             assert parameter > 0
     except AssertionError:
-        print('Parameter must be greater than 0')
+        print("Parameter must be greater than 0")
     try:
         assert num_pts_per_orbit > 0
     except AssertionError:
-        print('num_pts_per_orbit must be greater than 0')
+        print("num_pts_per_orbit must be greater than 0")
 
-    x = np.zeros((
-                    num_classes,
-                    num_orbits,
-                    num_pts_per_orbit,
-                    2
-                ))
+    x = np.zeros((num_classes, num_orbits, num_pts_per_orbit, 2))
 
     # generate dataset
     for cidx, p in enumerate(parameters):  # type: ignore
@@ -398,34 +439,35 @@ def generate_orbit_parallel(
             x_cur = x[cidx, :, i - 1, 0]
             y_cur = x[cidx, :, i - 1, 1]
 
-            x[cidx, :, i, 0] = (x_cur + p * y_cur * (1. - y_cur)) % 1
+            x[cidx, :, i, 0] = (x_cur + p * y_cur * (1.0 - y_cur)) % 1
             x_next = x[cidx, :, i, 0]
-            x[cidx, :, i, 1] = (y_cur + p * x_next * (1. - x_next)) % 1
+            x[cidx, :, i, 1] = (y_cur + p * x_next * (1.0 - x_next)) % 1
 
-    assert(not np.allclose(x[0, 0], x[0, 1]))
+    assert not np.allclose(x[0, 0], x[0, 1])
     return x
 
+
 def create_pd_orbits(
     orbits,
     num_classes,
-    homology_dimensions = (0, 1),
-    n_jobs = multiprocessing.cpu_count(),
-    ) -> torch.Tensor:
-    """ Computes the weak alpha persistence of the orbit data clouds.
+    homology_dimensions=(0, 1),
+    n_jobs=multiprocessing.cpu_count(),
+) -> torch.Tensor:
+    """Computes the weak alpha persistence of the orbit data clouds.
 
     Args:
-        orbits (np.array): 
+        orbits (np.array):
             Orbits of shape [n_points, 2]
-        homology_dimensions (tuple, optional): 
+        homology_dimensions (tuple, optional):
             Dimensions to compute the
             persistence diagrams.
             Defaults to (0, 1).
-        n_jobs (int, optional): 
+        n_jobs (int, optional):
             Number of cpus to use for parallel computation.
             Defaults to multiprocessing.cpu_count().
 
     Returns:
-        np.array: 
+        np.array:
             Array of persistence diagrams of shape
             [num_classes, num_orbits, num_persistence_points, 3].
             In the last dimension the first two values are the coordinates of
@@ -433,62 +475,58 @@ def create_pd_orbits(
             homology dimension.
     """
     # compute weak alpha persistence
-    wap = WeakAlphaPersistence(
-                        homology_dimensions=homology_dimensions,
-                        n_jobs=n_jobs
-                        )
+    wap = WeakAlphaPersistence(homology_dimensions=homology_dimensions, n_jobs=n_jobs)
     # c: class, o: orbit, p: point, d: dimension
-    orbits_stack = rearrange(orbits, 'c o p d -> (c o) p d')  # stack classes
+    orbits_stack = rearrange(orbits, "c o p d -> (c o) p d")  # stack classes
     diagrams = wap.fit_transform(orbits_stack)
     # shape: (num_classes * n_samples, n_features, 3)
 
     # combine class and orbit dimensions
     diagrams = rearrange(
-                            diagrams,
-                            '(c o) p d -> c o p d',
-                            c=num_classes  # type: ignore
-                        )
+        diagrams, "(c o) p d -> c o p d", c=num_classes  # type: ignore
+    )
     return diagrams
 
+
 def convert_pd_orbits_to_tensor(
     diagrams: np.ndarray,
     num_classes: int,
     num_orbits: int,
     num_homology_dimensions: int = 2,
-    ):
+):
     """[summary]
 
     Args:
-        diagrams (np.ndarray): 
+        diagrams (np.ndarray):
             [description]
-        num_classes (int): 
+        num_classes (int):
             [description]
-        num_orbits (int): 
+        num_orbits (int):
             [description]
-        num_homology_dimensions (int, optional): 
+        num_homology_dimensions (int, optional):
             [description]. Defaults to 2.
 
     Returns:
-        torch.tensor: 
+        torch.tensor:
             tuple of
             1.
             2. label tensor.
     """
     # c: class, o: orbit, p: point in persistence diagram,
     # d: coordinates + homology dimension
-    x = rearrange(
-                    diagrams,
-                    'c o p d -> (c o) p d',
-                    c=num_classes  # type: ignore
-                )
+    x = rearrange(diagrams, "c o p d -> (c o) p d", c=num_classes)  # type: ignore
     # convert homology dimension to one-hot encoding
     x = np.concatenate(
         (
             x[:, :, :2],
-            (np.eye(num_homology_dimensions)  # type: ignore
-             [x[:, :, -1].astype(np.int32)]),
+            (
+                np.eye(num_homology_dimensions)[  # type: ignore
+                    x[:, :, -1].astype(np.int32)
+                ]
+            ),
         ),
-        axis=-1)
+        axis=-1,
+    )
     # convert from [orbit, sequence_length, feature] to
     # [orbit, feature, sequence_length] to fit to the
     # input_shape of `SmallSetTransformer`
@@ -506,4 +544,4 @@ def convert_pd_orbits_to_tensor(
     x_tensor = torch.Tensor(x)
     y_tensor = torch.Tensor(y).long()
 
-    return x_tensor, y_tensor
\ No newline at end of file
+    return x_tensor, y_tensor
diff --git a/gdeep/data/datasets/persistence_diagrams_from_files.py b/gdeep/data/datasets/persistence_diagrams_from_files.py
index 895ad89b..17eacb68 100644
--- a/gdeep/data/datasets/persistence_diagrams_from_files.py
+++ b/gdeep/data/datasets/persistence_diagrams_from_files.py
@@ -5,7 +5,9 @@
 import torch
 from torch.utils.data import DataLoader, Dataset
 
-from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import OneHotEncodedPersistenceDiagram
+from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import (
+    OneHotEncodedPersistenceDiagram,
+)
 
 
 class PersistenceDiagramFromFiles(Dataset[Tuple[OneHotEncodedPersistenceDiagram, int]]):
@@ -13,24 +15,25 @@ class PersistenceDiagramFromFiles(Dataset[Tuple[OneHotEncodedPersistenceDiagram,
     lenght: int
     labels: Dict[int, int]
 
-    def __init__(self,
-                 file_path: str,
-                 ):
+    def __init__(
+        self,
+        file_path: str,
+    ):
         """
         Args:
-            file_path: 
+            file_path:
                 The path to the persistence diagrams.
                 num_homology_types: The number of homology types.
                 transform: A function that transforms the data.
         """
         self.file_path = file_path
-        self.len = len(os.listdir(join(self.file_path, 'diagrams')))
+        self.len = len(os.listdir(join(self.file_path, "diagrams")))
         self._load_labels()
 
     def __len__(self) -> int:
         """
         Return the length of the dataset.
-        
+
         Returns:
             The length of the dataset.
         """
@@ -39,33 +42,33 @@ def __len__(self) -> int:
     def _load_labels(self) -> None:
         """
         Load the labels of the dataset.
-        
+
         Returns:
             The labels of the dataset.
         """
         self.labels: Dict[int, int] = {}
         # load labels.csv with columns=["graph_idx", "label"] skipping the header
-        with open(join(self.file_path, 'labels.csv')) as f:
+        with open(join(self.file_path, "labels.csv")) as f:
             for line in f:
-                if line.startswith('graph_idx'):
+                if line.startswith("graph_idx"):
                     continue
-                graph_idx, label = line.strip().split(',')
+                graph_idx, label = line.strip().split(",")
                 self.labels[int(graph_idx)] = int(label)
         self.labels
 
     def __getitem__(self, index: int) -> Tuple[OneHotEncodedPersistenceDiagram, int]:
         """
         Return the item at the specified index.
-        
+
         Args:
-            index: 
+            index:
                 The index of the item.
-            
+
         Returns:
             The item at the specified index.
         """
-        diagram = OneHotEncodedPersistenceDiagram.load(join(self.file_path, 'diagrams',
-                                                            f'{index}.npy'))
+        diagram = OneHotEncodedPersistenceDiagram.load(
+            join(self.file_path, "diagrams", f"{index}.npy")
+        )
         label = self.labels[index]
         return diagram, label
-
diff --git a/gdeep/data/datasets/persistence_diagrams_from_graphs_builder.py b/gdeep/data/datasets/persistence_diagrams_from_graphs_builder.py
index 5ad55a52..0a596572 100644
--- a/gdeep/data/datasets/persistence_diagrams_from_graphs_builder.py
+++ b/gdeep/data/datasets/persistence_diagrams_from_graphs_builder.py
@@ -3,8 +3,9 @@
 
 import numpy as np
 import pandas as pd
-from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import \
-    OneHotEncodedPersistenceDiagram
+from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import (
+    OneHotEncodedPersistenceDiagram,
+)
 from gdeep.utility.extended_persistence import graph_extended_persistence_hks
 from gdeep.utility.constants import DEFAULT_GRAPH_DIR
 from torch_geometric.datasets import TUDataset  # type: ignore
@@ -19,14 +20,16 @@ class PersistenceDiagramFromGraphBuilder:
     This class is used to load the persistence diagrams of the graphs in a
     dataset. All graph datasets in the TUDataset class are supported.
     """
-    def __init__(self,
-                 dataset_name: str,
-                 diffusion_parameter: float,
-                 root: str = DEFAULT_GRAPH_DIR,
-                 ):
+
+    def __init__(
+        self,
+        dataset_name: str,
+        diffusion_parameter: float,
+        root: str = DEFAULT_GRAPH_DIR,
+    ):
         """
         Initialize the dataset.
-        
+
         Args:
             dataset_name:
                 The name of the graph dataset to load, e.g. "MUTAG".
@@ -38,31 +41,34 @@ def __init__(self,
         self.diffusion_parameter: float = diffusion_parameter
         self.num_homology_types: int = 4
         self.root: str = root
-        self.output_dir: str = os.path.join(root,
-                                       dataset_name + "_" + 
-                                       str(diffusion_parameter) +
-                                       "_extended_persistence")
-        
+        self.output_dir: str = os.path.join(
+            root,
+            dataset_name + "_" + str(diffusion_parameter) + "_extended_persistence",
+        )
+
     def create(self):
         # Check if the dataset exists in the specified directory
         if not os.path.exists(self.output_dir):
             print(f"Dataset {self.dataset_name} does not exist!")
             self._preprocess()
         else:
-            print(f"Dataset {self.dataset_name} already exists!"
-                " Skipping dataset will not be created.")
-        
+            print(
+                f"Dataset {self.dataset_name} already exists!"
+                " Skipping dataset will not be created."
+            )
+
     def __repr__(self) -> str:
         """
         Return a string representation of the dataset.
-        
+
         Returns:
             A string representation of the dataset.
         """
-        return f"{self.__class__.__name__}(dataset_name={self.dataset_name}, " \
-                f"diffusion_parameter={self.diffusion_parameter}, " \
-                f"root={self.root})"
-
+        return (
+            f"{self.__class__.__name__}(dataset_name={self.dataset_name}, "
+            f"diffusion_parameter={self.diffusion_parameter}, "
+            f"root={self.root})"
+        )
 
     def _preprocess(self) -> None:
         """
@@ -72,17 +78,19 @@ def _preprocess(self) -> None:
         method and then each diagram is saved in a separate npy file in the
         diagrams subdirectory of the output directory.
         The labels are saved in a csv file in the output directory.
-        
+
         Args:
             output_dir:
                 The directory where to save the persistence diagrams.
-        """        
+        """
         # Load the dataset
-        self.graph_dataset = TUDataset(root=DEFAULT_GRAPH_DIR,
-                                    name=self.dataset_name,
-                                    use_node_attr=False,
-                                    use_edge_attr=False)
-    
+        self.graph_dataset = TUDataset(
+            root=DEFAULT_GRAPH_DIR,
+            name=self.dataset_name,
+            use_node_attr=False,
+            use_edge_attr=False,
+        )
+
         # Create the directory if it does not exist
         if not os.path.exists(self.output_dir):
             os.makedirs(self.output_dir)
@@ -90,37 +98,35 @@ def _preprocess(self) -> None:
         else:
             # This should not be reached
             raise ValueError("Output directory already exists!")
-            
+
         num_graphs = len(self.graph_dataset)
-        
+
         labels: List[Tuple[int, int]] = []
-        
+
         print("Computing the persistence diagrams...")
-        for graph_idx, graph in tqdm(enumerate(self.graph_dataset),  # type: ignore
-                                     total=num_graphs):
-            
+        for graph_idx, graph in tqdm(
+            enumerate(self.graph_dataset), total=num_graphs  # type: ignore
+        ):
+
             # Get the adjacency matrix
             adj_mat: np.ndarray = to_dense_adj(graph.edge_index)[0].numpy()
-            
+
             # Compute the extended persistence
-            persistence_diagram_one_hot = \
-                graph_extended_persistence_hks(adj_mat, 
-                                               diffusion_parameter=
-                                               self.diffusion_parameter)
+            persistence_diagram_one_hot = graph_extended_persistence_hks(
+                adj_mat, diffusion_parameter=self.diffusion_parameter
+            )
             # Sort the diagram by the persistence lifetime, i.e. the second
             # column minus the first column
-            
+
             # Save the persistence diagram in a file
             persistence_diagram_one_hot.save(
-                os.path.join(self.output_dir, "diagrams",
-                              f"{graph_idx}.npy")
+                os.path.join(self.output_dir, "diagrams", f"{graph_idx}.npy")
             )
-            
+
             # Save the label
             labels.append((graph_idx, graph.y.item()))
 
         # Save the labels in a csv file
         pd.DataFrame(labels, columns=["graph_idx", "label"]).to_csv(
-            os.path.join(self.output_dir, "labels.csv"),
-            index=False
-            )
+            os.path.join(self.output_dir, "labels.csv"), index=False
+        )
diff --git a/gdeep/data/datasets/tori.py b/gdeep/data/datasets/tori.py
index 34a7df73..bc0e8a23 100644
--- a/gdeep/data/datasets/tori.py
+++ b/gdeep/data/datasets/tori.py
@@ -11,32 +11,26 @@
 Array = np.ndarray
 
 
-class Rotation():
-    """Class for rotations
-    """
-    def __init__(self, axis_0: int,
-                 axis_1: int,
-                 angle: float) -> None:
+class Rotation:
+    """Class for rotations"""
+
+    def __init__(self, axis_0: int, axis_1: int, angle: float) -> None:
         self._axis_0 = axis_0
         self._axis_1 = axis_1
         self._angle = angle
 
     def return_axis(self, idx: int):
-        return eval('self._axis_'+str(idx))
+        return eval("self._axis_" + str(idx))
 
     def return_angle(self) -> float:
         return self._angle
 
     def rotation_matrix(self) -> Array:
         rotation_matrix = np.identity(3)
-        rotation_matrix[self._axis_0, self._axis_0]\
-            = math.cos(self._angle)
-        rotation_matrix[self._axis_1, self._axis_1]\
-            = math.cos(self._angle)
-        rotation_matrix[self._axis_1, self._axis_0]\
-            = math.sin(self._angle)
-        rotation_matrix[self._axis_0, self._axis_1]\
-            = -math.sin(self._angle)
+        rotation_matrix[self._axis_0, self._axis_0] = math.cos(self._angle)
+        rotation_matrix[self._axis_1, self._axis_1] = math.cos(self._angle)
+        rotation_matrix[self._axis_1, self._axis_0] = math.sin(self._angle)
+        rotation_matrix[self._axis_0, self._axis_1] = -math.sin(self._angle)
         return rotation_matrix
 
 
@@ -48,6 +42,7 @@ class ToriDataset(Dataset[Tuple[Tensor, Tensor]]):
         name:
             name of the torus dataset to generate
     """
+
     def __init__(self, name: str, **kwargs) -> None:
         self.name = name
         if self.name == "DoubleTori":
@@ -62,18 +57,22 @@ def __getitem__(self, idx: int) -> Tuple[Tensor, Tensor]:
         It extracts the i-the element
         """
 
-        return self.out_tuple[0][idx], torch.tensor(self.out_tuple[1][idx]).to(torch.long)
+        return self.out_tuple[0][idx], torch.tensor(self.out_tuple[1][idx]).to(
+            torch.long
+        )
 
     def __len__(self) -> int:
         return len(self.out_tuple[0])
 
     @staticmethod
-    def _make_torus_point_cloud(label: int, 
-                                n_points: int,
-                                noise: float,
-                                rotation: Rotation,
-                                base_point: Array,
-                                radius: float=1.) -> Tuple[Array, Array]:
+    def _make_torus_point_cloud(
+        label: int,
+        n_points: int,
+        noise: float,
+        rotation: Rotation,
+        base_point: Array,
+        radius: float = 1.0,
+    ) -> Tuple[Array, Array]:
         """Generate point cloud of a single torus
 
         Args:
@@ -94,21 +93,23 @@ def _make_torus_point_cloud(label: int,
             (np.array, np.array):
                 data_points, labels
         """
-        torus_point_clouds = \
-            np.asarray(
+        torus_point_clouds = np.asarray(
+            [
                 [
-                    [
-                     (2 + radius*np.cos(s)) * np.cos(t) + noise *
-                     (np.random.rand(1)[0] - 0.5),
-                     (2 + radius*np.cos(s)) * np.sin(t) + noise *
-                     (np.random.rand(1)[0] - 0.5),
-                     radius*np.sin(s) + noise *
-                     (np.random.rand(1)[0] - 0.5),
-                    ]
-                    for t in range(n_points)
-                    for s in range(n_points)])
-        
-        torus_point_clouds = np.einsum("ij,kj->ki",  rotation.rotation_matrix(), torus_point_clouds)
+                    (2 + radius * np.cos(s)) * np.cos(t)
+                    + noise * (np.random.rand(1)[0] - 0.5),
+                    (2 + radius * np.cos(s)) * np.sin(t)
+                    + noise * (np.random.rand(1)[0] - 0.5),
+                    radius * np.sin(s) + noise * (np.random.rand(1)[0] - 0.5),
+                ]
+                for t in range(n_points)
+                for s in range(n_points)
+            ]
+        )
+
+        torus_point_clouds = np.einsum(
+            "ij,kj->ki", rotation.rotation_matrix(), torus_point_clouds
+        )
 
         torus_point_clouds += base_point
 
@@ -118,9 +119,15 @@ def _make_torus_point_cloud(label: int,
         return torus_point_clouds, torus_labels
 
     @staticmethod
-    def _make_torus_point_cloud2(label: int, n_points: int, noise: float,
-                                rotation: Rotation, base_point: Union[Array, List],
-                                radius1: float = 1., radius2: float = 1.) -> Tuple[Array, Array]:
+    def _make_torus_point_cloud2(
+        label: int,
+        n_points: int,
+        noise: float,
+        rotation: Rotation,
+        base_point: Union[Array, List],
+        radius1: float = 1.0,
+        radius2: float = 1.0,
+    ) -> Tuple[Array, Array]:
         """Generate point cloud of a single torus using
         2 radii for its definition
 
@@ -147,29 +154,31 @@ def _make_torus_point_cloud2(label: int, n_points: int, noise: float,
         torus_point_clouds = np.asarray(
             [
                 [
-                    (radius1 + radius2 * np.cos(s)) * np.cos(t) +
-                    noise * (np.random.rand(1)[0] - 0.5),
-                    (radius1 + radius2 * np.cos(s)) * np.sin(t) +
-                    noise * (np.random.rand(1)[0] - 0.5),
-                    radius2 * np.sin(s) + noise *
-                    (np.random.rand(1)[0] - 0.5),
+                    (radius1 + radius2 * np.cos(s)) * np.cos(t)
+                    + noise * (np.random.rand(1)[0] - 0.5),
+                    (radius1 + radius2 * np.cos(s)) * np.sin(t)
+                    + noise * (np.random.rand(1)[0] - 0.5),
+                    radius2 * np.sin(s) + noise * (np.random.rand(1)[0] - 0.5),
                 ]
                 for t in range(n_points)
-                for s in range(n_points)])
+                for s in range(n_points)
+            ]
+        )
 
-        torus_point_clouds = np.einsum("ij,kj->ki",
-                                       rotation.rotation_matrix(),
-                                       torus_point_clouds)
+        torus_point_clouds = np.einsum(
+            "ij,kj->ki", rotation.rotation_matrix(), torus_point_clouds
+        )
 
         torus_point_clouds += base_point
 
         # label tori with 2
-        torus_labels = label * np.ones(n_points ** 2)
+        torus_labels = label * np.ones(n_points**2)
 
         return torus_point_clouds, torus_labels
 
-    def _make_two_tori_dataset(self, entangled: bool = True,
-                               n_pts: int = 10) -> Tuple[Tensor, Array]:
+    def _make_two_tori_dataset(
+        self, entangled: bool = True, n_pts: int = 10
+    ) -> Tuple[Tensor, Array]:
         """Generates pandas Dataframe of two tori in 3D. The labels correspond to
         the different Tori.
 
@@ -184,40 +193,39 @@ def _make_two_tori_dataset(self, entangled: bool = True,
         torus_point_cloud = {}
         torus_labels = {}
         if entangled:
-            torus_point_cloud[0], torus_labels[0] = \
-                self._make_torus_point_cloud(0, n_pts, 0.0,
-                                             Rotation(1, 2,
-                                                      math.pi/2),
-                                             np.array([[0, 0, 0]]),
-                                             radius=.3)
-            torus_point_cloud[1], torus_labels[1] = \
-                self._make_torus_point_cloud(1, n_pts, 0.0,
-                                             Rotation(1, 2, 0),
-                                             np.array([[2, 0, 0]]),
-                                             radius=.3)
+            torus_point_cloud[0], torus_labels[0] = self._make_torus_point_cloud(
+                0,
+                n_pts,
+                0.0,
+                Rotation(1, 2, math.pi / 2),
+                np.array([[0, 0, 0]]),
+                radius=0.3,
+            )
+            torus_point_cloud[1], torus_labels[1] = self._make_torus_point_cloud(
+                1, n_pts, 0.0, Rotation(1, 2, 0), np.array([[2, 0, 0]]), radius=0.3
+            )
         else:
-            torus_point_cloud[0], torus_labels[0] = \
-                self._make_torus_point_cloud(0, n_pts, 0.0,
-                                             Rotation(1, 2,
-                                                      math.pi/2),
-                                             np.array([[0, 0, 0]]),
-                                             radius=.3)
-            torus_point_cloud[1], torus_labels[1] = \
-                self._make_torus_point_cloud(1, n_pts, 0.0,
-                                             Rotation(1, 2, 0),
-                                             np.array([[6, 0, 0]]),
-                                             radius=.3)
-
-        tori_point_cloud = np.concatenate((torus_point_cloud[0],
-                                           torus_point_cloud[1]),
-                                          axis=0)
-        tori_labels = np.concatenate((torus_labels[0],
-                                      torus_labels[1]),
-                                     axis=0)
+            torus_point_cloud[0], torus_labels[0] = self._make_torus_point_cloud(
+                0,
+                n_pts,
+                0.0,
+                Rotation(1, 2, math.pi / 2),
+                np.array([[0, 0, 0]]),
+                radius=0.3,
+            )
+            torus_point_cloud[1], torus_labels[1] = self._make_torus_point_cloud(
+                1, n_pts, 0.0, Rotation(1, 2, 0), np.array([[6, 0, 0]]), radius=0.3
+            )
+
+        tori_point_cloud = np.concatenate(
+            (torus_point_cloud[0], torus_point_cloud[1]), axis=0
+        )
+        tori_labels = np.concatenate((torus_labels[0], torus_labels[1]), axis=0)
         return torch.from_numpy(tori_point_cloud).to(torch.float32), tori_labels
 
-    def _make_entangled_tori_dataset(self, m: int = 2,
-                                     n_pts: int = 10) -> Tuple[Tensor, Array]:
+    def _make_entangled_tori_dataset(
+        self, m: int = 2, n_pts: int = 10
+    ) -> Tuple[Tensor, Array]:
         """Generates pandas Dataframe of m x m x m tori in 3D.
         The labels correspond to the different Tori.
 
@@ -231,15 +239,19 @@ def _make_entangled_tori_dataset(self, m: int = 2,
             (tuple):
                 the tuple for data and labels
         """
-        data1, labels1 = self._make_torus_point_cloud2(0, n_pts,
-                                                       0.1, Rotation(0, 1, 0),
-                                                       [0, 0, 0], 2., .5)
-        data2, labels2 = self._make_torus_point_cloud2(1, n_pts, 0.1,
-                                                       Rotation(1, 2, 180),
-                                                       [2, 0, 0], 2., .5)
-
-        translations = [[i*10, j*10, k*10] for i in range(m) for j
-                        in range(m) for k in range(m)]
+        data1, labels1 = self._make_torus_point_cloud2(
+            0, n_pts, 0.1, Rotation(0, 1, 0), [0, 0, 0], 2.0, 0.5
+        )
+        data2, labels2 = self._make_torus_point_cloud2(
+            1, n_pts, 0.1, Rotation(1, 2, 180), [2, 0, 0], 2.0, 0.5
+        )
+
+        translations = [
+            [i * 10, j * 10, k * 10]
+            for i in range(m)
+            for j in range(m)
+            for k in range(m)
+        ]
 
         tori_entangled = np.append(data1, data2, axis=0)
         labels = np.append(labels1, labels2)
@@ -267,6 +279,7 @@ def _make_blobs(m: int = 3, n_pts: int = 200) -> Tuple[Tensor, Array]:
             (tuple):
                 the tuple for data and labels
         """
-        data, lab = make_blobs(n_samples=n_pts, centers=m,
-                               n_features=3, random_state=42)
+        data, lab = make_blobs(
+            n_samples=n_pts, centers=m, n_features=3, random_state=42
+        )
         return torch.from_numpy(data).to(torch.float32), lab
diff --git a/gdeep/data/persistence_diagrams/__init__.py b/gdeep/data/persistence_diagrams/__init__.py
index 96109eca..526984fb 100644
--- a/gdeep/data/persistence_diagrams/__init__.py
+++ b/gdeep/data/persistence_diagrams/__init__.py
@@ -2,12 +2,12 @@
     OneHotEncodedPersistenceDiagram,
     get_one_hot_encoded_persistence_diagram_from_gtda,
     collate_fn_persistence_diagrams,
-    get_one_hot_encoded_persistence_diagram_from_gudhi_extended
+    get_one_hot_encoded_persistence_diagram_from_gudhi_extended,
 )
 
 __all__ = [
     "OneHotEncodedPersistenceDiagram",
     "collate_fn_persistence_diagrams",
     "get_one_hot_encoded_persistence_diagram_from_gtda",
-    "get_one_hot_encoded_persistence_diagram_from_gudhi_extended"
-]
\ No newline at end of file
+    "get_one_hot_encoded_persistence_diagram_from_gudhi_extended",
+]
diff --git a/gdeep/data/persistence_diagrams/one_hot_persistence_diagram.py b/gdeep/data/persistence_diagrams/one_hot_persistence_diagram.py
index e3a4f90d..bbf9a7e6 100644
--- a/gdeep/data/persistence_diagrams/one_hot_persistence_diagram.py
+++ b/gdeep/data/persistence_diagrams/one_hot_persistence_diagram.py
@@ -12,11 +12,12 @@
 
 class OneHotEncodedPersistenceDiagram:
     """This class represents a single one-hot encoded persistence diagram.
-    
+
     Args:
         data:
             The data of the persistence diagram.
     """
+
     _data: Tensor
     _homology_dimension_names: List[str]
 
@@ -26,49 +27,48 @@ def __init__(self, data: Tensor, homology_dimension_names: Optional[List[str]] =
         if homology_dimension_names is not None:
             self._homology_dimension_names = homology_dimension_names
         else:
-            self._homology_dimension_names = ["H_" + str(i) for i in range(self.get_num_homology_dimensions())]
-        assert len(self._homology_dimension_names) == self.get_num_homology_dimensions(), \
-            "The number of names must be equal to the number of homology dimensions."
+            self._homology_dimension_names = [
+                "H_" + str(i) for i in range(self.get_num_homology_dimensions())
+            ]
+        assert (
+            len(self._homology_dimension_names) == self.get_num_homology_dimensions()
+        ), "The number of names must be equal to the number of homology dimensions."
 
     def set_homology_dimension_names(self, homology_dimension_names: List[str]) -> None:
-        """This method sets the homology dimension names.
-        """
-        assert len(homology_dimension_names) == self.get_num_homology_dimensions(), \
-            "The number of names must be equal to the number of homology dimensions."
+        """This method sets the homology dimension names."""
+        assert (
+            len(homology_dimension_names) == self.get_num_homology_dimensions()
+        ), "The number of names must be equal to the number of homology dimensions."
         self._homology_dimension_names = homology_dimension_names
 
     def get_num_homology_dimensions(self) -> int:
-        """This method returns the number of homology dimensions.
-        """
+        """This method returns the number of homology dimensions."""
         return self._data.shape[-1] - 2
 
     def get_num_points(self) -> int:
-        """This method returns the number of points.
-        """
+        """This method returns the number of points."""
         return self._data.shape[0]
 
     def __repr__(self):
-        return (f"OneHotEncodedPersistenceDiagram({self._data.shape})\n"
-                f"{self.get_num_homology_dimensions()} homology dimensions\n"
-                f"{self._data.shape[0]} points\ndata:\n {self._data.__repr__()}")
+        return (
+            f"OneHotEncodedPersistenceDiagram({self._data.shape})\n"
+            f"{self.get_num_homology_dimensions()} homology dimensions\n"
+            f"{self._data.shape[0]} points\ndata:\n {self._data.__repr__()}"
+        )
 
     def save(self, path: str) -> None:
-        """This method saves the persistence diagram to a file.
-        """
+        """This method saves the persistence diagram to a file."""
         torch.save(self._data, path)
 
-    def get_points_in_homology_dimension(self, homology_dimension: int) -> \
-            'OneHotEncodedPersistenceDiagram':
-        """This method returns all points in a given homology dimension.
-        """
-        assert 0 <= homology_dimension < self.get_num_homology_dimensions(), \
-            "The homology dimension must be smaller than the number of homology dimensions."
+    def get_points_in_homology_dimension(
+        self, homology_dimension: int
+    ) -> "OneHotEncodedPersistenceDiagram":
+        """This method returns all points in a given homology dimension."""
+        assert (
+            0 <= homology_dimension < self.get_num_homology_dimensions()
+        ), "The homology dimension must be smaller than the number of homology dimensions."
         return OneHotEncodedPersistenceDiagram(
-            self._data[
-                torch.where(
-                    self._data[:, 2 + homology_dimension] == 1.0
-                )
-            ]
+            self._data[torch.where(self._data[:, 2 + homology_dimension] == 1.0)]
         )
 
     def get_raw_data(self) -> Tensor:
@@ -78,32 +78,28 @@ def get_raw_data(self) -> Tensor:
         return self._data
 
     @staticmethod
-    def load(path: str) -> 'OneHotEncodedPersistenceDiagram':
-        """This method loads a persistence diagram from a file.
-        """
+    def load(path: str) -> "OneHotEncodedPersistenceDiagram":
+        """This method loads a persistence diagram from a file."""
         return OneHotEncodedPersistenceDiagram(torch.load(path))
 
-    def get_all_points_in_homology_dimension(self, homology_dimension: int) -> \
-            'OneHotEncodedPersistenceDiagram':
-        """This method returns all points in a given homology dimension.
-        """
-        assert homology_dimension < self.get_num_homology_dimensions(), \
-            "The homology dimension must be smaller than the number of homology dimensions."
+    def get_all_points_in_homology_dimension(
+        self, homology_dimension: int
+    ) -> "OneHotEncodedPersistenceDiagram":
+        """This method returns all points in a given homology dimension."""
+        assert (
+            homology_dimension < self.get_num_homology_dimensions()
+        ), "The homology dimension must be smaller than the number of homology dimensions."
         return OneHotEncodedPersistenceDiagram(
-            self._data[
-                torch.where(
-                    self._data[:, 2 + homology_dimension] == 1.0
-                )
-            ]
+            self._data[torch.where(self._data[:, 2 + homology_dimension] == 1.0)]
         )
 
     def plot(self, names: Optional[List[str]] = None) -> gobj.Figure:  # type: ignore
         """This method plots the persistence diagram.
-        
+
         Args:
-            names: 
+            names:
                 The names of the homology dimensions.
-        
+
         Examples::
 
             pd =  torch.tensor\
@@ -127,38 +123,43 @@ def plot(self, names: Optional[List[str]] = None) -> gobj.Figure:  # type: ignor
 
         """
         names = self._homology_dimension_names if names is None else names
-        assert len(names) == self.get_num_homology_dimensions(), \
-            "The number of names must be equal to the number of homology dimensions."
+        assert (
+            len(names) == self.get_num_homology_dimensions()
+        ), "The number of names must be equal to the number of homology dimensions."
         pd = self._data.detach().numpy()
 
         # convert one-hot encoding to categorical encoding
         pd_categorical = pd[:, 2:].argmax(axis=1)
 
-        fig = _plot_diagram(np.concatenate([pd[:, :2], pd_categorical.reshape(-1, 1)], axis=1),
-                            names=names)
+        fig = _plot_diagram(
+            np.concatenate([pd[:, :2], pd_categorical.reshape(-1, 1)], axis=1),
+            names=names,
+        )
 
         return fig
 
-    def all_close(self, other: 'OneHotEncodedPersistenceDiagram',
-                  atol: float = 1e-7) -> bool:
-        """This method checks if the persistence diagrams are close.
-        """
+    def all_close(
+        self, other: "OneHotEncodedPersistenceDiagram", atol: float = 1e-7
+    ) -> bool:
+        """This method checks if the persistence diagrams are close."""
         for i in range(self.get_num_homology_dimensions()):
-            if not torch.allclose(self.get_all_points_in_homology_dimension(i).get_raw_data(),
-                                  other.get_all_points_in_homology_dimension(i).get_raw_data(),
-                                  atol=atol):
+            if not torch.allclose(
+                self.get_all_points_in_homology_dimension(i).get_raw_data(),
+                other.get_all_points_in_homology_dimension(i).get_raw_data(),
+                atol=atol,
+            ):
                 return False
         return True
 
     def get_lifetimes(self) -> Tensor:
-        """This method returns the lifetimes of the points.
-        """
+        """This method returns the lifetimes of the points."""
         return self._data[:, 1] - self._data[:, 0]
 
-    def filter_by_lifetime(self, min_lifetime: float, max_lifetime: float) -> \
-            'OneHotEncodedPersistenceDiagram':
+    def filter_by_lifetime(
+        self, min_lifetime: float, max_lifetime: float
+    ) -> "OneHotEncodedPersistenceDiagram":
         """This method filters the persistence diagram by lifetime.
-        
+
         Args:
             min_lifetime:
                 The minimum lifetime of the remaining points.
@@ -168,33 +169,33 @@ def filter_by_lifetime(self, min_lifetime: float, max_lifetime: float) -> \
         lifetime: Tensor = self.get_lifetimes()
         return OneHotEncodedPersistenceDiagram(
             self._data[  # type: ignore
-                torch.where(
-                    (lifetime >= min_lifetime) & (lifetime <= max_lifetime)
-                )
+                torch.where((lifetime >= min_lifetime) & (lifetime <= max_lifetime))
             ]
         )
 
     @staticmethod
-    def from_numpy(data: Array) -> 'OneHotEncodedPersistenceDiagram':
-        """This method creates a persistence diagram from a numpy array.
-        """
+    def from_numpy(data: Array) -> "OneHotEncodedPersistenceDiagram":
+        """This method creates a persistence diagram from a numpy array."""
         # assert data.dtype == np.float32, "The data must be of type np.float32. Otherwise,"\
         #     " the data will not be correctly converted."
-        return OneHotEncodedPersistenceDiagram(torch.from_numpy(data.astype(np.float32)))
+        return OneHotEncodedPersistenceDiagram(
+            torch.from_numpy(data.astype(np.float32))
+        )
 
 
-def collate_fn_persistence_diagrams(batch: List[Tuple[OneHotEncodedPersistenceDiagram, int]]) -> \
-        Tuple[List[Tensor], Tensor]:
+def collate_fn_persistence_diagrams(
+    batch: List[Tuple[OneHotEncodedPersistenceDiagram, int]]
+) -> Tuple[List[Tensor], Tensor]:
     """This function collates the data for the persistence diagram by padding the data, converting
     the data to tensors, converting the labels to tensors and generating masks for the valid
     entries.
-    
+
     The input is a list of tuples of the form (persistence diagram, label).
-    
+
     Args:
         batch:
             The list of tuples of the form (persistence diagram, label).
-            
+
     Returns:
         The data, the labels and the masks.
     """
@@ -203,8 +204,8 @@ def collate_fn_persistence_diagrams(batch: List[Tuple[OneHotEncodedPersistenceDi
     input_batch = torch.zeros(len(batch), max_num_points, num_homology_dimensions + 2)
     mask = torch.zeros(len(batch), max_num_points)
     for i, (pd, _) in enumerate(batch):
-        input_batch[i, :len(pd.get_raw_data()), :] = pd.get_raw_data()
-        mask[i, :len(pd.get_raw_data())] = 1.0
+        input_batch[i, : len(pd.get_raw_data()), :] = pd.get_raw_data()
+        mask[i, : len(pd.get_raw_data())] = 1.0
 
     label_batch = torch.tensor([x[1] for x in batch])
 
@@ -213,24 +214,25 @@ def collate_fn_persistence_diagrams(batch: List[Tuple[OneHotEncodedPersistenceDi
 
 def _check_if_valid(data) -> None:
     """This method checks if the data is a valid persistence diagram.
-    
+
     Args:
         data:
             The data to check.
-            
+
     Raises:
         ValueError:
             If the data is not a valid persistence diagram.
     """
 
     if data.ndimension() != 2:
-        raise ValueError("The input should be a 2-dimensional tensor."
-                         f"The input has {data.ndimension()} dimensions.")
-    assert data.shape[-1] > 2, \
-        "The input should have at least one homology dimensions."
-    assert torch.all(data[:, 2:] >= -1e-5) and \
-           torch.allclose(data[:, 2:].sum(dim=1), torch.tensor(1.0)), \
-           "The homology dimension should be one-hot encoded."
+        raise ValueError(
+            "The input should be a 2-dimensional tensor."
+            f"The input has {data.ndimension()} dimensions."
+        )
+    assert data.shape[-1] > 2, "The input should have at least one homology dimensions."
+    assert torch.all(data[:, 2:] >= -1e-5) and torch.allclose(
+        data[:, 2:].sum(dim=1), torch.tensor(1.0)
+    ), "The homology dimension should be one-hot encoded."
 
 
 # def _sort_by_lifetime(data: Tensor) -> Tensor:
@@ -240,55 +242,65 @@ def _check_if_valid(data) -> None:
 #         data[:, 1] - data[:, 0]
 #     ).argsort()]
 
-def get_one_hot_encoded_persistence_diagram_from_gtda(persistence_diagram: Array) \
-        -> OneHotEncodedPersistenceDiagram:
+
+def get_one_hot_encoded_persistence_diagram_from_gtda(
+    persistence_diagram: Array,
+) -> OneHotEncodedPersistenceDiagram:
     """This function takes a single persistence diagram from giotto-tda and returns a one-hot encoded
     persistence diagram.
-    
+
     Args:
         persistence_diagram:
             An array of shape (num_points, 3) where the first two columns
             represent the coordinates of the points and the third column
             represents the index of the homology dimension.
-        
+
     Returns:
         OneHotEncodedPersistenceDiagram:
             A one-hot encoded persistence diagram. If the persistence diagram has only one homology
             dimension, the third column will be filled with ones.
     """
-    assert persistence_diagram.ndim == 2 and persistence_diagram.shape[1] >= 2, \
-        "The input should be a 2-dimensional array of shape (num_points, 3) or (num_points, 2)."
+    assert (
+        persistence_diagram.ndim == 2 and persistence_diagram.shape[1] >= 2
+    ), "The input should be a 2-dimensional array of shape (num_points, 3) or (num_points, 2)."
 
     if persistence_diagram.shape[1] == 2:
         return OneHotEncodedPersistenceDiagram(
-            torch.stack((torch.tensor(persistence_diagram),
-                         torch.ones(persistence_diagram.shape[0])), dim=1))
+            torch.stack(
+                (
+                    torch.tensor(persistence_diagram),
+                    torch.ones(persistence_diagram.shape[0]),
+                ),
+                dim=1,
+            )
+        )
     else:
         homology_types: Set[int] = set([int(i) for i in persistence_diagram[:, 2]])
         type_to_one_hot_encoding: Dict[int, int] = {
             i: j for j, i in enumerate(homology_types)
         }
-        one_hot_encoding: Tensor = torch.zeros(persistence_diagram.shape[0],
-                                               len(homology_types))
+        one_hot_encoding: Tensor = torch.zeros(
+            persistence_diagram.shape[0], len(homology_types)
+        )
         # TODO: Fill the one-hot encoding in a vectorized manner.
         for i, j in enumerate(persistence_diagram[:, 2]):
             one_hot_encoding[i, type_to_one_hot_encoding[int(j)]] = 1
         birth_death_diagram: Tensor = torch.tensor(persistence_diagram[:, :2])
 
         return OneHotEncodedPersistenceDiagram(
-            torch.cat((birth_death_diagram,
-                         one_hot_encoding), dim=1).float())
+            torch.cat((birth_death_diagram, one_hot_encoding), dim=1).float()
+        )
 
 
 def get_one_hot_encoded_persistence_diagram_from_gudhi_extended(
-        diagram: Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]) \
-        -> OneHotEncodedPersistenceDiagram:
+    diagram: Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
+) -> OneHotEncodedPersistenceDiagram:
     """Convert an extended persistence diagram of a single graph to an
     array with one-hot encoded homology type.
     Args:
         diagram (Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]):
             The diagram of an extended persistence of a single graph.
-    
+
     Returns:
         np.ndarray:
             The diagram in one-hot encoded homology type of size
@@ -302,23 +314,21 @@ def get_one_hot_encoded_persistence_diagram_from_gudhi_extended(
 
     # One-hot encode the homology type
     homology_type: np.ndarray = np.array(
-        flatten_list_of_lists(
-            [[i] * lengths[i] for i in range(4)]
-        )
+        flatten_list_of_lists([[i] * lengths[i] for i in range(4)])
     )
     homology_type_one_hot = np.zeros((sum(lengths), 4))
     homology_type_one_hot[np.arange(sum(lengths)), homology_type] = 1
 
     # Concatenate the arrays
-    diagram_one_hot = np.concatenate([sub_diagram for sub_diagram in diagram],
-                                     axis=0)
-    diagram_one_hot = np.concatenate([diagram_one_hot, homology_type_one_hot],
-                                     axis=1)
+    diagram_one_hot = np.concatenate([sub_diagram for sub_diagram in diagram], axis=0)
+    diagram_one_hot = np.concatenate([diagram_one_hot, homology_type_one_hot], axis=1)
 
     return OneHotEncodedPersistenceDiagram.from_numpy(diagram_one_hot)
 
 
-def _plot_diagram(diagram, names: List[str], homology_dimensions=None, plotly_params=None):
+def _plot_diagram(
+    diagram, names: List[str], homology_dimensions=None, plotly_params=None
+):
     """Plot a single persistence diagram.
 
     Args:
@@ -361,14 +371,16 @@ def _plot_diagram(diagram, names: List[str], homology_dimensions=None, plotly_pa
     max_val_display = max_val + extra_space
 
     fig = gobj.Figure()
-    fig.add_trace(gobj.Scatter(
-        x=[min_val_display, max_val_display],
-        y=[min_val_display, max_val_display],
-        mode="lines",
-        line={"dash": "dash", "width": 1, "color": "black"},
-        showlegend=False,
-        hoverinfo="none"
-    ))
+    fig.add_trace(
+        gobj.Scatter(
+            x=[min_val_display, max_val_display],
+            y=[min_val_display, max_val_display],
+            mode="lines",
+            line={"dash": "dash", "width": 1, "color": "black"},
+            showlegend=False,
+            hoverinfo="none",
+        )
+    )
 
     for dim in homology_dimensions:
         name = names[int(dim)]
@@ -377,20 +389,27 @@ def _plot_diagram(diagram, names: List[str], homology_dimensions=None, plotly_pa
             subdiagram, axis=0, return_inverse=True, return_counts=True
         )
         hovertext = [
-            f"{tuple(unique[unique_row_index][:2])}" +
-            (
+            f"{tuple(unique[unique_row_index][:2])}"
+            + (
                 f", multiplicity: {counts[unique_row_index]}"
-                if counts[unique_row_index] > 1 else ""
+                if counts[unique_row_index] > 1
+                else ""
             )
             for unique_row_index in inverse
         ]
         y = subdiagram[:, 1]
         if has_posinfinite_death:
             y[np.isposinf(y)] = posinfinity_val  # type: ignore
-        fig.add_trace(gobj.Scatter(
-            x=subdiagram[:, 0], y=y, mode="markers",
-            hoverinfo="text", hovertext=hovertext, name=name
-        ))
+        fig.add_trace(
+            gobj.Scatter(
+                x=subdiagram[:, 0],
+                y=y,
+                mode="markers",
+                hoverinfo="text",
+                hovertext=hovertext,
+                name=name,
+            )
+        )
 
     fig.update_layout(
         width=500,
@@ -408,14 +427,16 @@ def _plot_diagram(diagram, names: List[str], homology_dimensions=None, plotly_pa
             "linecolor": "black",
             "mirror": False,
             "showexponent": "all",
-            "exponentformat": "e"
+            "exponentformat": "e",
         },
         yaxis1={
             "title": "Death",
             "side": "left",
             "type": "linear",
             "range": [min_val_display, max_val_display],
-            "autorange": False, "scaleanchor": "x", "scaleratio": 1,
+            "autorange": False,
+            "scaleanchor": "x",
+            "scaleratio": 1,
             "ticks": "outside",
             "showline": True,
             "zeroline": True,
@@ -423,22 +444,24 @@ def _plot_diagram(diagram, names: List[str], homology_dimensions=None, plotly_pa
             "linecolor": "black",
             "mirror": False,
             "showexponent": "all",
-            "exponentformat": "e"
+            "exponentformat": "e",
         },
-        plot_bgcolor="white"
+        plot_bgcolor="white",
     )
 
     # Add a horizontal dashed line for points with infinite death
     if has_posinfinite_death:
-        fig.add_trace(gobj.Scatter(
-            x=[min_val_display, max_val_display],
-            y=[posinfinity_val, posinfinity_val],  # type: ignore
-            mode="lines",
-            line={"dash": "dash", "width": 0.5, "color": "black"},
-            showlegend=True,
-            name=u"\u221E",
-            hoverinfo="none"
-        ))
+        fig.add_trace(
+            gobj.Scatter(
+                x=[min_val_display, max_val_display],
+                y=[posinfinity_val, posinfinity_val],  # type: ignore
+                mode="lines",
+                line={"dash": "dash", "width": 0.5, "color": "black"},
+                showlegend=True,
+                name="\u221E",
+                hoverinfo="none",
+            )
+        )
 
     # Update traces and layout according to user input
     if plotly_params:
diff --git a/gdeep/data/preprocessors/__init__.py b/gdeep/data/preprocessors/__init__.py
index 34c0ef46..3886dff7 100644
--- a/gdeep/data/preprocessors/__init__.py
+++ b/gdeep/data/preprocessors/__init__.py
@@ -1,22 +1,23 @@
-
 from .normalization import Normalization
 from .tokenizer_translation import TokenizerTranslation
 from .tokenizer_qa import TokenizerQA
 from .to_tensor_image import ToTensorImage
 from .tokenizer_text_classification import TokenizerTextClassification
 from .filter_persistence_diagram_by_lifetime import FilterPersistenceDiagramByLifetime
-from .filter_persistence_diagram_by_homology_dimension import FilterPersistenceDiagramByHomologyDimension
+from .filter_persistence_diagram_by_homology_dimension import (
+    FilterPersistenceDiagramByHomologyDimension,
+)
 from .normalization_persistence_diagram import NormalizationPersistenceDiagram
 from .min_max_normalization_persistence_diagram import MinMaxScalarPersistenceDiagram
 
 __all__ = [
-    'Normalization',
-    'TokenizerTranslation',
-    'TokenizerQA',
-    'ToTensorImage',
-    'TokenizerTextClassification',
-    'NormalizationPersistenceDiagram',
-    'FilterPersistenceDiagramByLifetime',
-    'FilterPersistenceDiagramByHomologyDimension',
-    'MinMaxScalarPersistenceDiagram',
-    ]
+    "Normalization",
+    "TokenizerTranslation",
+    "TokenizerQA",
+    "ToTensorImage",
+    "TokenizerTextClassification",
+    "NormalizationPersistenceDiagram",
+    "FilterPersistenceDiagramByLifetime",
+    "FilterPersistenceDiagramByHomologyDimension",
+    "MinMaxScalarPersistenceDiagram",
+]
diff --git a/gdeep/data/preprocessors/filter_persistence_diagram_by_homology_dimension.py b/gdeep/data/preprocessors/filter_persistence_diagram_by_homology_dimension.py
index ec210c07..f8ea8d4c 100644
--- a/gdeep/data/preprocessors/filter_persistence_diagram_by_homology_dimension.py
+++ b/gdeep/data/preprocessors/filter_persistence_diagram_by_homology_dimension.py
@@ -2,8 +2,9 @@
 from typing import Any, List, Tuple, TypeVar
 
 import torch
-from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import \
-    OneHotEncodedPersistenceDiagram
+from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import (
+    OneHotEncodedPersistenceDiagram,
+)
 from torch.utils.data import Dataset
 
 from ..abstract_preprocessing import AbstractPreprocessing
@@ -12,52 +13,56 @@
 Tensor = torch.Tensor
 PD = OneHotEncodedPersistenceDiagram
 
-T = TypeVar('T')
+T = TypeVar("T")
 
-class FilterPersistenceDiagramByHomologyDimension(AbstractPreprocessing[Tuple[PD, T], Tuple[PD, T]]):
+
+class FilterPersistenceDiagramByHomologyDimension(
+    AbstractPreprocessing[Tuple[PD, T], Tuple[PD, T]]
+):
     """This class filters the persistence diagrams of a dataset by their homology dimension.
-    
+
     Here we assume that the dataset is a tuple of (persistence diagram, label) and that
     the points in the diagram are sorted by ascending lifetime. This is an invariant of
     the OneHotEncodedPersistenceDiagram class but could go wrong if the diagrams are modified
     in a way that breaks this invariant.
-     
+
     Args:
         homology_dimensions_to_filter:
             The homology dimensions of the points in the diagram that should be kept.
     """
+
     is_fitted: bool
     homology_dimensions_to_filter_by: List[int]
-    
-    def __init__(self,
-                 homology_dimensions_to_filter: List[int]):
+
+    def __init__(self, homology_dimensions_to_filter: List[int]):
         self.homology_dimensions_to_filter_by = homology_dimensions_to_filter
         self.is_fitted = False
-        
+
     def fit_to_dataset(self, dataset: Dataset[Tuple[PD, T]]) -> None:
         """This method does nothing."""
         self.is_fitted = True
-        
+
     def __call__(self, item: Tuple[PD, T]) -> Tuple[PD, T]:
         """Filters the persistence diagram by homology dimension.
-        
+
         Args:
-            item: 
+            item:
                 A tuple of (persistence diagram, label).
-        
+
         Returns:
             A tuple of (persistence diagram, label).
         """
         if not self.is_fitted:
-            raise RuntimeError("The filter is not fitted to any dataset. "
-                               "Please call fit_to_dataset() first.")
+            raise RuntimeError(
+                "The filter is not fitted to any dataset. "
+                "Please call fit_to_dataset() first."
+            )
         out: List[Tensor] = []
         for homology_dimension in self.homology_dimensions_to_filter_by:
-            out.append(item[0].get_all_points_in_homology_dimension(homology_dimension).get_raw_data())
-        
-        return PD(torch.cat(out)), item[1]
-        
+            out.append(
+                item[0]
+                .get_all_points_in_homology_dimension(homology_dimension)
+                .get_raw_data()
+            )
 
-        
-        
-        
\ No newline at end of file
+        return PD(torch.cat(out)), item[1]
diff --git a/gdeep/data/preprocessors/filter_persistence_diagram_by_lifetime.py b/gdeep/data/preprocessors/filter_persistence_diagram_by_lifetime.py
index cc87791c..00f17cbe 100644
--- a/gdeep/data/preprocessors/filter_persistence_diagram_by_lifetime.py
+++ b/gdeep/data/preprocessors/filter_persistence_diagram_by_lifetime.py
@@ -2,8 +2,9 @@
 from typing import Any, Tuple, TypeVar
 
 import torch
-from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import \
-    OneHotEncodedPersistenceDiagram
+from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import (
+    OneHotEncodedPersistenceDiagram,
+)
 from torch.utils.data import Dataset
 
 from ..abstract_preprocessing import AbstractPreprocessing
@@ -12,55 +13,57 @@
 Tensor = torch.Tensor
 PD = OneHotEncodedPersistenceDiagram
 
-T = TypeVar('T')
+T = TypeVar("T")
 
 
-class FilterPersistenceDiagramByLifetime(AbstractPreprocessing[Tuple[PD, T], Tuple[PD, T]]):
+class FilterPersistenceDiagramByLifetime(
+    AbstractPreprocessing[Tuple[PD, T], Tuple[PD, T]]
+):
     """This class filters the persistence diagrams of a dataset by their lifetime, i.e.
     the difference between the birth and death coordinates.
-    
+
     Here we assume that the dataset is a tuple of (persistence diagram, label) and that
     the points in the diagram are sorted by ascending lifetime. This is an invariant of
     the OneHotEncodedPersistenceDiagram class but could go wrong if the diagrams are modified
     in a way that breaks this invariant.
-    
+
     Args:
         min_lifetime:
             The minimum lifetime of the points in the diagram.
         max_lifetime:
             The maximum lifetime of the points in the diagram.
     """
+
     is_fitted: bool
     max_lifetime: float
     min_lifetime: float
-    
-    def __init__(self,
-                 min_lifetime: float,
-                 max_lifetime: float):
+
+    def __init__(self, min_lifetime: float, max_lifetime: float):
         self.min_lifetime = min_lifetime
         self.max_lifetime = max_lifetime
         self.is_fitted = False
-        
+
     def fit_to_dataset(self, dataset: Dataset[Tuple[PD, T]]) -> None:
         """This method does nothing."""
         self.is_fitted = True
-        
+
     def __call__(self, item: Tuple[PD, T]) -> Tuple[PD, T]:
         """Filters the persistence diagram by its lifetime.
-        
+
         Args:
             item:
                 A tuple of (persistence diagram, label).
-        
+
         Returns:
             A tuple of (persistence diagram, label).
         """
         if not self.is_fitted:
-            raise RuntimeError("The filter is not fitted to any dataset. "
-                               "Please call fit_to_dataset() first.")
+            raise RuntimeError(
+                "The filter is not fitted to any dataset. "
+                "Please call fit_to_dataset() first."
+            )
         out: Tensor = item[0].get_raw_data()
         lifetime: Tensor = out[:, 1] - out[:, 0]
         mask: Tensor = (lifetime >= self.min_lifetime) & (lifetime <= self.max_lifetime)
         out = out[mask]  # type: ignore
         return (PD(out), item[1])
-    
\ No newline at end of file
diff --git a/gdeep/data/preprocessors/min_max_normalization_persistence_diagram.py b/gdeep/data/preprocessors/min_max_normalization_persistence_diagram.py
index af530301..0d8587b9 100644
--- a/gdeep/data/preprocessors/min_max_normalization_persistence_diagram.py
+++ b/gdeep/data/preprocessors/min_max_normalization_persistence_diagram.py
@@ -2,8 +2,9 @@
 from typing import Any, Tuple, TypeVar
 
 import torch
-from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import \
-    OneHotEncodedPersistenceDiagram
+from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import (
+    OneHotEncodedPersistenceDiagram,
+)
 from gdeep.utility import DEVICE
 from matplotlib.pyplot import axis
 from torch.utils.data import Dataset
@@ -15,34 +16,39 @@
 Tensor = torch.Tensor
 PD = OneHotEncodedPersistenceDiagram
 
-T = TypeVar('T')
+T = TypeVar("T")
 
 
-class MinMaxScalarPersistenceDiagram(AbstractPreprocessing[Tuple[Tensor, T], Tuple[Tensor, T]]):
-    """This class runs the standard min-max normalisation on the birth and death times 
+class MinMaxScalarPersistenceDiagram(
+    AbstractPreprocessing[Tuple[Tensor, T], Tuple[Tensor, T]]
+):
+    """This class runs the standard min-max normalisation on the birth and death times
     of the persistence diagrams. For example. The transformation is:
     X_scaled = X_std * (max - min) + min
     """
+
     is_fitted: bool
     min_max = Tuple[Tensor, Tensor]
-    
+
     def __init__(self):
         self.is_fitted = False
-        
+
     def fit_to_dataset(self, dataset: Dataset[Tuple[Tensor, T]]) -> None:
         self.min_max = _compute_min_max_of_dataset(
-            TransformingDataset(dataset, lambda x: (torch.cat([x[0].min(), x[0].max()]), x[1]))
+            TransformingDataset(
+                dataset, lambda x: (torch.cat([x[0].min(), x[0].max()]), x[1])
+            )
         )
         self.is_fitted = True
-        
+
     def __call__(self, item: Tuple[Tensor, T]) -> Tuple[Tensor, T]:
         if not self.is_fitted:
-            raise RuntimeError("The normalisation is not fitted to any dataset."
-                               " Please call fit_to_dataset() first.")
+            raise RuntimeError(
+                "The normalisation is not fitted to any dataset."
+                " Please call fit_to_dataset() first."
+            )
         if not torch.all(self.min_max[1] > 0):  # type: ignore
             warnings.warn("The max values are zeros! Adding 1e-7")
             self.min_max = (self.min_max[0], self.min_max[1] + 1e-7)  # type: ignore
-        out = (item[0] - self.min_max[0])/ self.min_max[1]  # type: ignore
+        out = (item[0] - self.min_max[0]) / self.min_max[1]  # type: ignore
         return (out, item[1])
-    
-    
diff --git a/gdeep/data/preprocessors/min_max_scalar.py b/gdeep/data/preprocessors/min_max_scalar.py
index e5cc11e5..c8549f99 100644
--- a/gdeep/data/preprocessors/min_max_scalar.py
+++ b/gdeep/data/preprocessors/min_max_scalar.py
@@ -2,8 +2,9 @@
 from typing import Any, Tuple, TypeVar
 
 import torch
-from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import \
-    OneHotEncodedPersistenceDiagram
+from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import (
+    OneHotEncodedPersistenceDiagram,
+)
 from gdeep.utility import DEVICE
 from matplotlib.pyplot import axis
 from torch.utils.data import Dataset
@@ -14,15 +15,16 @@
 Tensor = torch.Tensor
 PD = OneHotEncodedPersistenceDiagram
 
-T = TypeVar('T')
+T = TypeVar("T")
 
 
 class MinMaxScalarP(AbstractPreprocessing[Tuple[Tensor, Any], Tuple[Tensor, Any]]):
     """This class runs the standard min-max scaling on all elements of the dataset.
-    
+
     For example. The transformation is:
-    `` X_scaled = X_std * (max - min) + min `` 
+    `` X_scaled = X_std * (max - min) + min ``
     """
+
     is_fitted: bool
     min_max = Tuple[Tensor, Tensor]
 
@@ -30,15 +32,15 @@ def __init__(self):
         self.is_fitted = False
 
     def fit_to_dataset(self, dataset: Dataset[Tuple[Tensor, Any]]) -> None:
-        self.min_max = _compute_min_max_of_dataset(
-            dataset
-        )
+        self.min_max = _compute_min_max_of_dataset(dataset)
         self.is_fitted = True
 
     def __call__(self, item: Tuple[Tensor, Any]) -> Tuple[Tensor, Any]:
         if not self.is_fitted:
-            raise RuntimeError("The normalisation is not fitted to any dataset."
-                               " Please call fit_to_dataset() first.")
+            raise RuntimeError(
+                "The normalisation is not fitted to any dataset."
+                " Please call fit_to_dataset() first."
+            )
         if not torch.all(self.min_max[1] > 0):  # type: ignore
             warnings.warn("The max values are zeros! Adding 1e-7")
             self.min_max = (self.min_max[0], self.min_max[1] + 1e-7)  # type: ignore
@@ -46,11 +48,13 @@ def __call__(self, item: Tuple[Tensor, Any]) -> Tuple[Tensor, Any]:
         return item
 
 
-def _compute_min_max_of_dataset(dataset: Dataset[Tuple[Tensor, Any]]) -> Tuple[Tensor, Tensor]:
+def _compute_min_max_of_dataset(
+    dataset: Dataset[Tuple[Tensor, Any]]
+) -> Tuple[Tensor, Tensor]:
     """Compute the min and max of the whole dataset"""
     min: Tensor = dataset[0][0].clone().to(DEVICE)
     max: Tensor = dataset[0][0].clone().to(DEVICE)
-    for item in dataset:   # type: ignore
+    for item in dataset:  # type: ignore
         min = torch.min(min, item[0])
         max = torch.max(max, item[0])
     return min, max
diff --git a/gdeep/data/preprocessors/normalization.py b/gdeep/data/preprocessors/normalization.py
index 5dd431df..b106819a 100644
--- a/gdeep/data/preprocessors/normalization.py
+++ b/gdeep/data/preprocessors/normalization.py
@@ -5,7 +5,9 @@
 import torch
 from torch.utils.data import Dataset
 
-from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import OneHotEncodedPersistenceDiagram
+from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import (
+    OneHotEncodedPersistenceDiagram,
+)
 
 from ..abstract_preprocessing import AbstractPreprocessing
 from ..transforming_dataset import TransformingDataset
@@ -14,7 +16,7 @@
 Tensor = torch.Tensor
 PD = OneHotEncodedPersistenceDiagram
 
-T = TypeVar('T')
+T = TypeVar("T")
 
 
 class Normalization(AbstractPreprocessing[Tuple[Tensor, T], Tuple[Tensor, T]]):
@@ -23,6 +25,7 @@ class Normalization(AbstractPreprocessing[Tuple[Tensor, T], Tuple[Tensor, T]]):
     shape ``(C, H, W)``, the average will and the standard deviations
     will be tensors of shape ``(C, H, W)``
     """
+
     is_fitted: bool
     mean: Tensor
     stddev: Tensor
@@ -37,14 +40,17 @@ def fit_to_dataset(self, dataset: Dataset[Tuple[Tensor, T]]) -> None:
 
     def __call__(self, item: Tuple[Tensor, T]) -> Tuple[Tensor, T]:
         if not self.is_fitted:
-            raise RuntimeError("The normalisation is not fitted to any dataset."
-                               " Please call fit_to_dataset() first.")
+            raise RuntimeError(
+                "The normalisation is not fitted to any dataset."
+                " Please call fit_to_dataset() first."
+            )
         if not torch.all(self.stddev > 0):
             warnings.warn("The standard deviation contains zeros! Adding 1e-7")
             self.stddev = self.stddev + 1e-7
-        out = (item[0] - self.mean)/ self.stddev
+        out = (item[0] - self.mean) / self.stddev
         return (out, item[1])
 
+
 def _compute_mean_of_dataset(dataset: Dataset[Tuple[Tensor, Any]]) -> Tensor:
     """Compute the mean of the whole dataset"""
     mean: Tensor = torch.zeros(dataset[0][0].shape, dtype=torch.float64)
@@ -55,11 +61,17 @@ def _compute_mean_of_dataset(dataset: Dataset[Tuple[Tensor, Any]]) -> Tensor:
             mean = (mean * idx + dataset[idx][0]) / (idx + 1)
     return mean
 
-def _compute_stddev_of_dataset(dataset: Dataset[Tuple[Tensor, Any]], mean: Tensor) -> Tensor:
+
+def _compute_stddev_of_dataset(
+    dataset: Dataset[Tuple[Tensor, Any]], mean: Tensor
+) -> Tensor:
     """Compute the stddev of the whole dataset"""
+
     def square_diff_from_mean(x: Tuple[Tensor, Any]) -> Tuple[Tensor, Any]:
         return (x[0] - mean) ** 2, x[1]
-    mean_normalized_dataset = TransformingDataset[Tuple[Tensor, Any], Tuple[Tensor, Any]] \
-        (dataset, square_diff_from_mean)
+
+    mean_normalized_dataset = TransformingDataset[
+        Tuple[Tensor, Any], Tuple[Tensor, Any]
+    ](dataset, square_diff_from_mean)
     stddev: Tensor = _compute_mean_of_dataset(mean_normalized_dataset)
     return stddev.sqrt()
diff --git a/gdeep/data/preprocessors/normalization_persistence_diagram.py b/gdeep/data/preprocessors/normalization_persistence_diagram.py
index dd7b29bb..14d754c2 100644
--- a/gdeep/data/preprocessors/normalization_persistence_diagram.py
+++ b/gdeep/data/preprocessors/normalization_persistence_diagram.py
@@ -2,8 +2,9 @@
 from typing import Any, Tuple, TypeVar
 
 import torch
-from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import \
-    OneHotEncodedPersistenceDiagram
+from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import (
+    OneHotEncodedPersistenceDiagram,
+)
 from torch.utils.data import Dataset
 
 from ..abstract_preprocessing import AbstractPreprocessing
@@ -14,45 +15,55 @@
 Tensor = torch.Tensor
 PD = OneHotEncodedPersistenceDiagram
 
-T = TypeVar('T')
+T = TypeVar("T")
 
 
-class NormalizationPersistenceDiagram(AbstractPreprocessing[Tuple[PD, T], Tuple[PD, T]]):
+class NormalizationPersistenceDiagram(
+    AbstractPreprocessing[Tuple[PD, T], Tuple[PD, T]]
+):
     """This class runs the standard normalisation on the birth and death coordinates
     of the persistence diagrams of a dataset accross all the homology dimensions.
-    
+
     The one-hot encoded persistence diagrams are kept as is.
     """
+
     is_fitted: bool
     mean: Tensor
     stddev: Tensor
     num_homology_dimensions: int
 
-    def __init__(self,
-                 num_homology_dimensions: int):
+    def __init__(self, num_homology_dimensions: int):
         self.is_fitted = False
         self.num_homology_dimensions = num_homology_dimensions
 
     def fit_to_dataset(self, dataset: Dataset[Tuple[PD, T]]) -> None:
         # compute the mean and set the last entries to zero
         self.mean = _compute_mean_of_dataset(
-            TransformingDataset(dataset, lambda x: (x[0].get_raw_data().mean(dim=0), x[1]))
+            TransformingDataset(
+                dataset, lambda x: (x[0].get_raw_data().mean(dim=0), x[1])
             )
-        self.mean = self.mean * torch.tensor([1.0, 1.0] + [0.0] * (self.num_homology_dimensions))
+        )
+        self.mean = self.mean * torch.tensor(
+            [1.0, 1.0] + [0.0] * (self.num_homology_dimensions)
+        )
         self.stddev = _compute_mean_of_dataset(
-            TransformingDataset(dataset, lambda x: (((x[0].get_raw_data() - self.mean)**2).mean(dim=0), x[1]))
+            TransformingDataset(
+                dataset,
+                lambda x: (((x[0].get_raw_data() - self.mean) ** 2).mean(dim=0), x[1]),
             )
-        self.stddev = (self.stddev 
-                       * torch.tensor([1.0, 1.0] + [0.0] * (self.num_homology_dimensions))
-                       + torch.tensor([0.0, 0.0] + [1.0] * (self.num_homology_dimensions))
-                       )
+        )
+        self.stddev = self.stddev * torch.tensor(
+            [1.0, 1.0] + [0.0] * (self.num_homology_dimensions)
+        ) + torch.tensor([0.0, 0.0] + [1.0] * (self.num_homology_dimensions))
         self.stddev = torch.sqrt(self.stddev)
         self.is_fitted = True
 
     def __call__(self, item: Tuple[PD, T]) -> Tuple[PD, T]:
         if not self.is_fitted:
-            raise RuntimeError("The normalisation is not fitted to any dataset."
-                               " Please call fit_to_dataset() first.")
+            raise RuntimeError(
+                "The normalisation is not fitted to any dataset."
+                " Please call fit_to_dataset() first."
+            )
         if not torch.all(self.stddev > 0):
             warnings.warn("The standard deviation contains zeros! Adding 1e-7")
             self.stddev = self.stddev + 1e-7
diff --git a/gdeep/data/preprocessors/to_tensor_image.py b/gdeep/data/preprocessors/to_tensor_image.py
index 82653863..5b76c750 100644
--- a/gdeep/data/preprocessors/to_tensor_image.py
+++ b/gdeep/data/preprocessors/to_tensor_image.py
@@ -9,23 +9,24 @@
 # type definition
 Tensor = torch.Tensor
 
-class ToTensorImage(AbstractPreprocessing[Tuple[Any, Any],
-                                          Tuple[Tensor, Tensor]]):
+
+class ToTensorImage(AbstractPreprocessing[Tuple[Any, Any], Tuple[Tensor, Tensor]]):
     """Class to preprocess image files for classification
-      tasks
+    tasks
 
-          Args:
-              size :
-                  Desired output size. If size is a sequence like (h, w),
-                  output size will be matched to this. If size is an int,
-                  smaller edge of the image will be matched to this number.
-                  I.e, if height > width, then image will be rescaled to
-                  ``(size * height / width, size)``.
+        Args:
+            size :
+                Desired output size. If size is a sequence like (h, w),
+                output size will be matched to this. If size is an int,
+                smaller edge of the image will be matched to this number.
+                I.e, if height > width, then image will be rescaled to
+                ``(size * height / width, size)``.
 
     """
 
     is_fitted: bool
     size: Union[int, List[int]]
+
     def __init__(self, size: Union[int, List[int]]) -> None:
         self.size = size
         self.is_fitted = True
@@ -35,4 +36,4 @@ def fit_to_dataset(self, dataset: Dataset[Any]) -> None:
         pass
 
     def __call__(self, datum: Tuple[Any, Any]) -> Tuple[Tensor, Tensor]:
-        return ToTensor()(Resize(self.size)(datum[0])), torch.tensor(datum[1],dtype=torch.long)  # type: ignore
+        return ToTensor()(Resize(self.size)(datum[0])), torch.tensor(datum[1], dtype=torch.long)  # type: ignore
diff --git a/gdeep/data/preprocessors/tokenizer_qa.py b/gdeep/data/preprocessors/tokenizer_qa.py
index d459526e..52d97d9c 100644
--- a/gdeep/data/preprocessors/tokenizer_qa.py
+++ b/gdeep/data/preprocessors/tokenizer_qa.py
@@ -16,8 +16,12 @@
 Tensor = torch.Tensor
 
 
-class TokenizerQA(AbstractPreprocessing[Tuple[str, str, List[str], List[int]],
-                                        Tuple[Union[Tensor, List[Tensor]], Tensor]]):
+class TokenizerQA(
+    AbstractPreprocessing[
+        Tuple[str, str, List[str], List[int]],
+        Tuple[Union[Tensor, List[Tensor]], Tensor],
+    ]
+):
     """Class to preprocess text dataloaders for Q&A
     tasks. The type of dataset is assumed to be of the
     form ``(string,string,list[string], list[string])``.
@@ -41,6 +45,7 @@ class TokenizerQA(AbstractPreprocessing[Tuple[str, str, List[str], List[int]],
                                TokenizerQA())
 
     """
+
     is_fitted: bool
     max_length: int
     vocabulary: Optional[Vocab]
@@ -49,10 +54,11 @@ class TokenizerQA(AbstractPreprocessing[Tuple[str, str, List[str], List[int]],
     pad_item: int
     ordered_dict: OrderedDict
 
-    def __init__(self, vocabulary: Optional[Vocab] = None,
-                 tokenizer: Optional[partial] = None):
+    def __init__(
+        self, vocabulary: Optional[Vocab] = None, tokenizer: Optional[partial] = None
+    ):
         if tokenizer is None:
-            self.tokenizer = get_tokenizer('basic_english')
+            self.tokenizer = get_tokenizer("basic_english")
         else:
             self.tokenizer = tokenizer
         self.vocabulary = vocabulary
@@ -60,14 +66,16 @@ def __init__(self, vocabulary: Optional[Vocab] = None,
         self.max_length_context: int = 0
         self.is_fitted = False
 
-    def fit_to_dataset(self, dataset: Dataset[Tuple[str, str, List[str], List[int]]]) -> None:
+    def fit_to_dataset(
+        self, dataset: Dataset[Tuple[str, str, List[str], List[int]]]
+    ) -> None:
         """Method to fit the vocabulary to the input text
 
         Args:
             dataset:
                 the dataset to fit to
 
-            """
+        """
         self.counter = Counter()
         self.ordered_dict = OrderedDict()
         for (context, question, answer, init_position) in dataset:  # type: ignore
@@ -77,41 +85,62 @@ def fit_to_dataset(self, dataset: Dataset[Tuple[str, str, List[str], List[int]]]
             self.max_length_question = max(self.max_length_question, len(self.tokenizer(question)))  # type: ignore
 
         if self.vocabulary is None:
-            self.ordered_dict = OrderedDict(sorted(self.counter.items(),
-                                                   key=lambda x: x[1],
-                                                   reverse=True))
+            self.ordered_dict = OrderedDict(
+                sorted(self.counter.items(), key=lambda x: x[1], reverse=True)
+            )
             self.vocabulary = vocab(self.ordered_dict)
-            unk_token = '<unk>'  # type: ignore
+            unk_token = "<unk>"  # type: ignore
             if unk_token not in self.vocabulary:
                 self.vocabulary.insert_token(unk_token, 0)
             self.vocabulary.set_default_index(self.vocabulary[unk_token])
         self.pad_item = 0
         self.is_fitted = True
 
-    def __call__(self, datum: Tuple[str, str, List[str], List[int]]) -> Tuple[Union[Tensor, List[Tensor]], Tensor]:
+    def __call__(
+        self, datum: Tuple[str, str, List[str], List[int]]
+    ) -> Tuple[Union[Tensor, List[Tensor]], Tensor]:
         """This method implement the transformation once fitted."""
         # if not self.is_fitted:
         #    self.load_pretrained(".")
         if self.vocabulary:
-            text_pipeline: Callable[[str], List[int]] = lambda x: [self.vocabulary[token] for token in  # type: ignore
-                                                                   self.tokenizer(x)]  # type: ignore
+            text_pipeline: Callable[[str], List[int]] = lambda x: [
+                self.vocabulary[token] for token in self.tokenizer(x)  # type: ignore
+            ]  # type: ignore
         else:
-            raise MissingVocabularyError("Please fit this preprocessor to initialise the vocabulary")
-
-        processed_context = torch.tensor(text_pipeline(datum[0]),
-                                         dtype=torch.int64).to(DEVICE)
-        out_context = torch.cat([processed_context,
-                                 self.pad_item * torch.ones(self.max_length_context -
-                                                            processed_context.shape[0]).to(DEVICE)])
-        processed_question = torch.tensor(text_pipeline(datum[1]),
-                                          dtype=torch.int64).to(DEVICE)
-
-        out_question = torch.cat([processed_question,
-                                  self.pad_item * torch.ones(self.max_length_question -
-                                                             processed_question.shape[0]).to(DEVICE)])
+            raise MissingVocabularyError(
+                "Please fit this preprocessor to initialise the vocabulary"
+            )
+
+        processed_context = torch.tensor(text_pipeline(datum[0]), dtype=torch.int64).to(
+            DEVICE
+        )
+        out_context = torch.cat(
+            [
+                processed_context,
+                self.pad_item
+                * torch.ones(self.max_length_context - processed_context.shape[0]).to(
+                    DEVICE
+                ),
+            ]
+        )
+        processed_question = torch.tensor(
+            text_pipeline(datum[1]), dtype=torch.int64
+        ).to(DEVICE)
+
+        out_question = torch.cat(
+            [
+                processed_question,
+                self.pad_item
+                * torch.ones(self.max_length_question - processed_question.shape[0]).to(
+                    DEVICE
+                ),
+            ]
+        )
 
         pos_init_char = datum[3][0]
         pos_init = len(self.tokenizer(datum[0][:pos_init_char]))  # type: ignore
         pos_end = pos_init + len(self.tokenizer(datum[2][0]))  # type: ignore
         out_list = [out_context.to(torch.long), out_question.to(torch.long)]
-        return out_list, torch.stack((torch.tensor(pos_init), torch.tensor(pos_end))).to(torch.long)
+        return out_list, torch.stack(
+            (torch.tensor(pos_init), torch.tensor(pos_end))
+        ).to(torch.long)
diff --git a/gdeep/data/preprocessors/tokenizer_text_classification.py b/gdeep/data/preprocessors/tokenizer_text_classification.py
index 765ed01e..a97c8b23 100644
--- a/gdeep/data/preprocessors/tokenizer_text_classification.py
+++ b/gdeep/data/preprocessors/tokenizer_text_classification.py
@@ -1,6 +1,5 @@
 from collections import Counter, OrderedDict
-from typing import Callable, Tuple, \
-    Union, Any, Optional, List, Dict
+from typing import Callable, Tuple, Union, Any, Optional, List, Dict
 from functools import partial as Partial  # noqa
 
 import torch
@@ -17,8 +16,9 @@
 Tensor = torch.Tensor
 
 
-class TokenizerTextClassification(AbstractPreprocessing[Tuple[Any, str],
-                                                        Tuple[Tensor, Tensor]]):
+class TokenizerTextClassification(
+    AbstractPreprocessing[Tuple[Any, str], Tuple[Tensor, Tensor]]
+):
     """Preprocessing class. This class is useful to convert the data format
     ``(label, text)`` into the proper tensor format ``( word_embedding, label)``.
     The labels shouzld be interagers; f they ar string, they will be converted.
@@ -39,10 +39,11 @@ class TokenizerTextClassification(AbstractPreprocessing[Tuple[Any, str],
     counter: Dict[str, int]
     counter_label: Dict[str, int]
 
-    def __init__(self, tokenizer: Optional[Partial] = None,
-                 vocabulary: Optional[Vocab] = None):
+    def __init__(
+        self, tokenizer: Optional[Partial] = None, vocabulary: Optional[Vocab] = None
+    ):
         if tokenizer is None:
-            self.tokenizer = get_tokenizer('basic_english')
+            self.tokenizer = get_tokenizer("basic_english")
         else:
             self.tokenizer = tokenizer
 
@@ -72,12 +73,12 @@ def fit_to_dataset(self, dataset: Dataset[Tuple[Any, str]]) -> None:
             self.max_length = max(self.max_length, len(self.tokenizer(text)))  # type: ignore
         # build the vocabulary
         if not self.vocabulary:
-            self.ordered_dict = OrderedDict(sorted(self.counter.items(),
-                                                   key=lambda x: x[1],
-                                                   reverse=True))
+            self.ordered_dict = OrderedDict(
+                sorted(self.counter.items(), key=lambda x: x[1], reverse=True)
+            )
             self.list_of_possible_labels = list(self.counter_label.keys())
             self.vocabulary = vocab(self.ordered_dict)
-            unk_token = '<unk>'  # type: ignore
+            unk_token = "<unk>"  # type: ignore
             if unk_token not in self.vocabulary:
                 self.vocabulary.insert_token(unk_token, 0)
             self.vocabulary.set_default_index(self.vocabulary[unk_token])
@@ -96,24 +97,32 @@ def __call__(self, datum: Tuple[Any, str]) -> Tuple[Tensor, Tensor]:
         # if not self.is_fitted:
         #    self.load_pretrained(".")
         if self.vocabulary:
-            text_pipeline: Callable[[str], List[int]] = lambda x: [self.vocabulary[token]  # type: ignore
-                                                                   for token in self.tokenizer(x)]  # type: ignore
+            text_pipeline: Callable[[str], List[int]] = lambda x: [
+                self.vocabulary[token] for token in self.tokenizer(x)  # type: ignore
+            ]  # type: ignore
         else:
-            raise MissingVocabularyError("Please fit this preprocessor to initialise the vocabulary")
+            raise MissingVocabularyError(
+                "Please fit this preprocessor to initialise the vocabulary"
+            )
         pad_item: int = 0
 
         _text = datum[1]
         # if isinstance(_text, tuple) or isinstance(_text, list):
         #    _text = _text[1]
-        processed_text = torch.tensor(text_pipeline(_text),
-                                      dtype=torch.long).to(DEVICE)
+        processed_text = torch.tensor(text_pipeline(_text), dtype=torch.long).to(DEVICE)
         # convert to tensors (padded)
-        out_text = torch.cat([processed_text,
-                              pad_item * torch.ones(self.max_length - processed_text.shape[0]
-                                                    ).to(DEVICE)]).to(torch.long)
+        out_text = torch.cat(
+            [
+                processed_text,
+                pad_item
+                * torch.ones(self.max_length - processed_text.shape[0]).to(DEVICE),
+            ]
+        ).to(torch.long)
         # preprocess labels
         label_pipeline = lambda x: torch.tensor(x, dtype=torch.long)
 
         _label = datum[0]
-        out_label = label_pipeline(self.list_of_possible_labels.index(_label)).to(DEVICE)
+        out_label = label_pipeline(self.list_of_possible_labels.index(_label)).to(
+            DEVICE
+        )
         return out_text, out_label
diff --git a/gdeep/data/preprocessors/tokenizer_translation.py b/gdeep/data/preprocessors/tokenizer_translation.py
index 3a6cd8ef..322f2488 100644
--- a/gdeep/data/preprocessors/tokenizer_translation.py
+++ b/gdeep/data/preprocessors/tokenizer_translation.py
@@ -1,6 +1,5 @@
 from collections import Counter, OrderedDict
-from typing import Callable, Tuple, \
-    Optional, List, Dict, Union
+from typing import Callable, Tuple, Optional, List, Dict, Union
 
 import torch
 from torchtext.data.utils import get_tokenizer  # type: ignore
@@ -14,9 +13,9 @@
 Tensor = torch.Tensor
 
 
-class TokenizerTranslation(AbstractPreprocessing[Tuple[str, str],
-                                                 Tuple[Union[Tensor, List[Tensor]], Tensor]
-                           ]):
+class TokenizerTranslation(
+    AbstractPreprocessing[Tuple[str, str], Tuple[Union[Tensor, List[Tensor]], Tensor]]
+):
     """Class to preprocess text dataloaders for translation
     tasks. The Dataset type is supposed to be ``(string, string)``.
     The padding item is supposed to be of index 0.
@@ -47,7 +46,8 @@ class TokenizerTranslation(AbstractPreprocessing[Tuple[str, str],
         textds = TransformingDataset(ds_tr,
             TokenizerTranslation())
 
-        """
+    """
+
     if_fitted: bool
     vocabulary: Optional[Dict[str, int]]
     vocabulary_target: Optional[Dict[str, int]]
@@ -56,17 +56,20 @@ class TokenizerTranslation(AbstractPreprocessing[Tuple[str, str],
     ordered_dict: Dict[str, int]
     ordered_dict_target: Dict[str, int]
 
-    def __init__(self, vocabulary: Optional[Dict[str, int]] = None,
-                 vocabulary_target: Optional[Dict[str, int]] = None,
-                 tokenizer: Optional[Callable[[str], List[str]]] = None,
-                 tokenizer_target: Optional[Callable[[str], List[str]]] = None):
+    def __init__(
+        self,
+        vocabulary: Optional[Dict[str, int]] = None,
+        vocabulary_target: Optional[Dict[str, int]] = None,
+        tokenizer: Optional[Callable[[str], List[str]]] = None,
+        tokenizer_target: Optional[Callable[[str], List[str]]] = None,
+    ):
 
         if tokenizer is None:
-            self.tokenizer = get_tokenizer('basic_english')
+            self.tokenizer = get_tokenizer("basic_english")
         else:
             self.tokenizer = tokenizer
         if tokenizer_target is None:
-            self.tokenizer_target = get_tokenizer('basic_english')
+            self.tokenizer_target = get_tokenizer("basic_english")
         else:
             self.tokenizer_target = tokenizer_target
         self.vocabulary = vocabulary
@@ -83,30 +86,38 @@ def fit_to_dataset(self, data):
         for text in data:
             self.counter.update(self.tokenizer(text[0]))
             self.counter_target.update(self.tokenizer_target(text[1]))
-            self.max_length_source = max(self.max_length_source, len(self.tokenizer(text[0])))
-            self.max_length_target = max(self.max_length_target, len(self.tokenizer_target(text[1])))
+            self.max_length_source = max(
+                self.max_length_source, len(self.tokenizer(text[0]))
+            )
+            self.max_length_target = max(
+                self.max_length_target, len(self.tokenizer_target(text[1]))
+            )
         # self.vocabulary = Vocab(counter, min_freq=1)
         if self.vocabulary is None:
-            self.ordered_dict = OrderedDict(sorted(self.counter.items(),
-                                                   key=lambda x: x[1],
-                                                   reverse=True))
+            self.ordered_dict = OrderedDict(
+                sorted(self.counter.items(), key=lambda x: x[1], reverse=True)
+            )
             self.vocabulary = vocab(self.ordered_dict)
-            unk_token = '<unk>'  # type: ignore
-            if unk_token not in self.vocabulary: self.vocabulary.insert_token(unk_token, 0)
+            unk_token = "<unk>"  # type: ignore
+            if unk_token not in self.vocabulary:
+                self.vocabulary.insert_token(unk_token, 0)
             self.vocabulary.set_default_index(self.vocabulary[unk_token])
         if self.vocabulary_target is None:
-            self.ordered_dict_target = OrderedDict(sorted(self.counter_target.items(),
-                                                          key=lambda x: x[1],
-                                                          reverse=True))
+            self.ordered_dict_target = OrderedDict(
+                sorted(self.counter_target.items(), key=lambda x: x[1], reverse=True)
+            )
             self.vocabulary_target = vocab(self.ordered_dict_target)
-            unk_token = '<unk>'  # type: ignore
-            if unk_token not in self.vocabulary_target: self.vocabulary_target.insert_token(unk_token, 0)
+            unk_token = "<unk>"  # type: ignore
+            if unk_token not in self.vocabulary_target:
+                self.vocabulary_target.insert_token(unk_token, 0)
             self.vocabulary_target.set_default_index(self.vocabulary_target[unk_token])
             # self.vocabulary_target = Vocab(self.counter_target)
         self.is_fitted = True
         # self.save_pretrained(".")
 
-    def __call__(self, datum: Tuple[str, str]) -> Tuple[Union[List[Tensor], Tensor], Tensor]:
+    def __call__(
+        self, datum: Tuple[str, str]
+    ) -> Tuple[Union[List[Tensor], Tensor], Tensor]:
         """This method is applied to each item of
         the dataset and
         transforms it following the rule described
@@ -120,31 +131,51 @@ def __call__(self, datum: Tuple[str, str]) -> Tuple[Union[List[Tensor], Tensor],
         # if not self.is_fitted:
         #    self.load_pretrained(".")
         if self.vocabulary:
-            text_pipeline: Callable[[str], List[int]] = lambda x: [self.vocabulary[token] for token in  # type: ignore
-                                                                   self.tokenizer(x)]  # type: ignore
+            text_pipeline: Callable[[str], List[int]] = lambda x: [
+                self.vocabulary[token] for token in self.tokenizer(x)  # type: ignore
+            ]  # type: ignore
         else:
-            raise MissingVocabularyError("Please fit this preprocessor to initialise the vocabulary")
+            raise MissingVocabularyError(
+                "Please fit this preprocessor to initialise the vocabulary"
+            )
 
         if self.vocabulary_target:
-            text_pipeline_target: Callable[[str], List[int]] = \
-                lambda xx: [self.vocabulary_target[token] for token in self.tokenizer_target(xx)]  # type: ignore
+            text_pipeline_target: Callable[[str], List[int]] = lambda xx: [
+                self.vocabulary_target[token] for token in self.tokenizer_target(xx)
+            ]  # type: ignore
         else:
-            raise MissingVocabularyError("Please fit this preprocessor to initialise the vocabulary")
+            raise MissingVocabularyError(
+                "Please fit this preprocessor to initialise the vocabulary"
+            )
 
         pad_item: int = 0
         pad_item_target: int = 0
 
-        processed_text = torch.tensor(text_pipeline(datum[0]),
-                                      dtype=torch.long).to(DEVICE)
-        processed_text_target = torch.tensor(text_pipeline_target(datum[1]),
-                                             dtype=torch.long).to(DEVICE)
+        processed_text = torch.tensor(text_pipeline(datum[0]), dtype=torch.long).to(
+            DEVICE
+        )
+        processed_text_target = torch.tensor(
+            text_pipeline_target(datum[1]), dtype=torch.long
+        ).to(DEVICE)
         # convert to tensors (padded)
-        out_source = torch.cat([processed_text,
-                                pad_item * torch.ones(self.max_length_source - processed_text.shape[0]
-                                                      ).to(DEVICE)]).to(torch.long)
-        out_target = torch.cat([processed_text_target,
-                                pad_item_target * torch.ones(self.max_length_target - processed_text_target.shape[0]
-                                                             ).to(DEVICE)]).to(torch.long)
+        out_source = torch.cat(
+            [
+                processed_text,
+                pad_item
+                * torch.ones(self.max_length_source - processed_text.shape[0]).to(
+                    DEVICE
+                ),
+            ]
+        ).to(torch.long)
+        out_target = torch.cat(
+            [
+                processed_text_target,
+                pad_item_target
+                * torch.ones(
+                    self.max_length_target - processed_text_target.shape[0]
+                ).to(DEVICE),
+            ]
+        ).to(torch.long)
         x = [out_source.to(torch.long), out_target.to(torch.long)]
         y = out_target.to(torch.long)  # .clone().detach()
         return x, y
diff --git a/gdeep/models/extractor.py b/gdeep/models/extractor.py
index 9678289c..e8e40148 100644
--- a/gdeep/models/extractor.py
+++ b/gdeep/models/extractor.py
@@ -36,14 +36,15 @@ class ModelExtractor:
     """
 
     def __init__(
-            self, model: torch.nn.Module, loss_fn: Callable[[Tensor, Tensor], Tensor]
+        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 _send_to_device(self, x: Union[Tensor, List[Tensor]], y: Tensor) \
-            -> Tuple[Tensor, Union[Tensor, List[Tensor]], Tensor]:
+    def _send_to_device(
+        self, x: Union[Tensor, List[Tensor]], y: Tensor
+    ) -> Tuple[Tensor, Union[Tensor, List[Tensor]], Tensor]:
         """use this private method to send the
         ``x`` and ``y`` to the ``DEVICE``.
 
@@ -84,7 +85,7 @@ def _send_to_device(self, x: Union[Tensor, List[Tensor]], y: Tensor) \
         return prediction, x, y
 
     def get_decision_boundary(
-            self, input_example: Tensor, n_epochs: int = 100, precision: float = 0.1
+        self, input_example: Tensor, n_epochs: int = 100, precision: float = 0.1
     ) -> Tensor:
         """Compute the decision boundary for self.model
         with self.loss_fn
@@ -148,7 +149,9 @@ def get_activations(self, x: Union[Tensor, List[Tensor]]) -> List[Tensor]:
             handle = layer.register_forward_hook(saved_output_layers)
             hook_handles.append(handle)
         self.model.eval()
-        _, x, _ = self._send_to_device(x, torch.rand((1, 1)))  # here we also run the model(x)
+        _, x, _ = self._send_to_device(
+            x, torch.rand((1, 1))
+        )  # here we also run the model(x)
 
         for handle in hook_handles:
             handle.remove()
@@ -189,8 +192,9 @@ def get_layers_grads(self) -> List[Tensor]:
             output.append(v.grad)
         return output
 
-    def get_gradients(self, batch: Tuple[Union[Tensor, List[Tensor]], Tensor],
-                      **kwargs) -> Tuple[List[Tensor], List[Tensor]]:
+    def get_gradients(
+        self, batch: Tuple[Union[Tensor, List[Tensor]], Tensor], **kwargs
+    ) -> Tuple[List[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=`
diff --git a/gdeep/models/simple_nn.py b/gdeep/models/simple_nn.py
index 111ab081..677a27f1 100644
--- a/gdeep/models/simple_nn.py
+++ b/gdeep/models/simple_nn.py
@@ -14,20 +14,21 @@ class FFNet(nn.Module):
             the last one which has ``sigmoid`` as activation function.
             The first number is the dimension of the input! The final
             of the output
-        
+
         activation (callable, default = torch.nn.functional.relu):
             this is the activation function that will be applied between
             each layer of the fully connected network
     """
-    
-    def __init__(self, arch=(2, 3, 3, 2), activation = F.relu):        
+
+    def __init__(self, arch=(2, 3, 3, 2), activation=F.relu):
         super(FFNet, self).__init__()
         self.activation = activation
-        self.linears = nn.ModuleList([nn.Linear(arch[i], arch[i+1]) for i in range(len(arch)-1)])
-        
+        self.linears = nn.ModuleList(
+            [nn.Linear(arch[i], arch[i + 1]) for i in range(len(arch) - 1)]
+        )
 
     def forward(self, x):
         for i, l in enumerate(self.linears[:-1]):
             x = self.activation(l(x))
         x = self.linears[-1](x)
-        return x
\ No newline at end of file
+        return x
diff --git a/gdeep/models/utils.py b/gdeep/models/utils.py
index 08a4a3f8..e0f4a0c7 100644
--- a/gdeep/models/utils.py
+++ b/gdeep/models/utils.py
@@ -8,7 +8,7 @@
 
 
 class SaveOutput(ABC):
-    """ General class for saving outputs of
+    """General class for saving outputs of
     neural networks.
     Outputs will be stored in a list 'outputs'
     """
@@ -32,7 +32,7 @@ def get_outputs(self) -> List[Any]:
 
 
 class SaveNodeOutput(SaveOutput):
-    """ Class for saving activations of a node in
+    """Class for saving activations of a node in
     a neural network.
 
     Args:
diff --git a/gdeep/search/benchmark.py b/gdeep/search/benchmark.py
index 4526a85a..0e6981f5 100644
--- a/gdeep/search/benchmark.py
+++ b/gdeep/search/benchmark.py
@@ -125,8 +125,8 @@ def forward(self, x):
     def __init__(
         self,
         models_dicts: List[Dict[str, Union[torch.nn.Module, str]]],
-        dataloaders_dicts: List[Dict[
-            str, Union[List[DataLoader[Tuple[Tensor, Tensor]]], str]]
+        dataloaders_dicts: List[
+            Dict[str, Union[List[DataLoader[Tuple[Tensor, Tensor]]], str]]
         ],
         loss_fn: Callable[[Tensor, Tensor], Tensor],
         writer: SummaryWriter,
@@ -173,7 +173,7 @@ def start(
         writer_tag: str = "",
     ) -> None:
         """Method to be called when starting the benchmarking
-        
+
         Args:
             optimizer:
                 a torch optimizers class (not the instance)
@@ -242,7 +242,7 @@ def _inner_function(
     ) -> None:
         """private method to run the inner
         function of the benchmark loops
-        
+
         Args:
             model:
                 dictionary defining the model name
@@ -276,7 +276,7 @@ def _benchmarking_param(
 ) -> None:
     """Function to be used as pseudo-decorator for
     benchmarking loops
-    
+
     Args:
         fun (Callable):
             the function to decorate
diff --git a/gdeep/search/hpo.py b/gdeep/search/hpo.py
index a5d39d52..0f241528 100644
--- a/gdeep/search/hpo.py
+++ b/gdeep/search/hpo.py
@@ -134,7 +134,7 @@ class HyperParameterOptimization(Trainer):
         pruner :
             Instance of an optuna pruner, can be user-defined
         sampler :
-            If left unspecified, ``TPESample`` is used during single-objective 
+            If left unspecified, ``TPESample`` is used during single-objective
             optimization and ``NSGAIISampler`` during multi-objective optimization
         db_url:
             name of the database to connect to. For example
@@ -142,7 +142,7 @@ class HyperParameterOptimization(Trainer):
         study_name:
             name of the optuna study
     Examples::
-    
+
         from gdeep.search import HyperParameterOptimization
         # initialise hpo, you need a `trainer`!
         search = HyperParameterOptimization(trainer, "accuracy", 2, best_not_last=True)
@@ -163,6 +163,7 @@ class HyperParameterOptimization(Trainer):
             n_accumulated_grads=2,
         )
     """
+
     is_pipe: bool
     df_res: pd.DataFrame
     study: Study
@@ -179,7 +180,7 @@ def __init__(
         study_name: Optional[str] = None,
     ):
         self.best_not_last_gs = best_not_last
-        self.best_val_acc_gs = 0.
+        self.best_val_acc_gs = 0.0
         self.best_val_loss_gs = np.inf
         self.list_res: List[Any] = []
         self.db_url = db_url
@@ -222,11 +223,11 @@ def _initialise_new_model(
     ) -> torch.nn.Module:
         """private method to find the maximal compatible set
         between models and hyperparameters
-        
+
         Args:
             models_hyperparam (dict):
                 model selected hyperparameters
-        
+
         Returns:
             nn.Module
                 torch nn.Module
@@ -263,7 +264,7 @@ def _initialise_new_model(
 
     def _objective(self, trial: BaseTrial, config: HPOConfig):
         """default callback function for optuna's study
-        
+
         Args:
             trial:
                 the independent variable
@@ -385,12 +386,12 @@ def _extract_model_and_dataset_name(self, writer_tag: str) -> Tuple[str, str]:
         if index_ds == -1:
             dataset_name = self.pipe.dataloaders[0].dataset.__class__.__name__
         else:
-            dataset_name = writer_tag[index_ds + 8: writer_tag.find("|Model:")]
+            dataset_name = writer_tag[index_ds + 8 : writer_tag.find("|Model:")]
         index_md = writer_tag.find("|Model:")
         if index_md == -1:
             model_name = self.pipe.model.__class__.__name__
         else:
-            model_name = writer_tag[index_md + 7: writer_tag.find("/")]
+            model_name = writer_tag[index_md + 7 : writer_tag.find("/")]
         return model_name, dataset_name
 
     def start(
@@ -509,7 +510,8 @@ def _inner_optimisat_fun(
         self,
         model: Union[torch.nn.Module, Dict[str, torch.nn.Module]],
         dataloaders: Union[
-            List[DataLoader[Tuple[Union[Tensor, List[Tensor]], Tensor]]], Dict[str, DataLoader]
+            List[DataLoader[Tuple[Union[Tensor, List[Tensor]], Tensor]]],
+            Dict[str, DataLoader],
         ],
         config: HPOConfig,
     ) -> None:
@@ -601,7 +603,7 @@ def _store_to_list_each_step(
     ) -> List[Any]:
         """Private method to store all the HPO parameters
         of one trial
-        
+
         Args:
             trial (optuna.trial):
                 the trial at hand
@@ -618,7 +620,7 @@ def _store_to_list_each_step(
         Returns:
             list:
                 list of HPOs and metrics. the first element is the
-                run name and the last two are the metrics (loss 
+                run name and the last two are the metrics (loss
                 and accuracy)
         """
         if list_res is None:
@@ -888,11 +890,17 @@ def _new_suggest_categorical(trial, name: str, choices: Sequence) -> Any:
         """
 
         if isinstance(choices, list) or isinstance(choices, tuple):
-            if isinstance(choices[0], str) or isinstance(choices[0], int) or isinstance(choices[0], float):
+            if (
+                isinstance(choices[0], str)
+                or isinstance(choices[0], int)
+                or isinstance(choices[0], float)
+            ):
                 return trial.suggest_categorical(name, choices)
             else:  # in case optuna cannot handle the types
                 dict_choices = {x.__name__: x for x in choices}
-                key = trial.suggest_categorical(name, dict_choices.keys())  # random choice on the names
+                key = trial.suggest_categorical(
+                    name, dict_choices.keys()
+                )  # random choice on the names
                 return dict_choices[key]  # return the corresponding value
         return random.choice(choices)
 
diff --git a/gdeep/topology_layers/__init__.py b/gdeep/topology_layers/__init__.py
index 8cd19551..3144ff05 100644
--- a/gdeep/topology_layers/__init__.py
+++ b/gdeep/topology_layers/__init__.py
@@ -5,9 +5,9 @@
 from .persformer_wrapper import PersformerWrapper
 
 __all__ = [
-    'PersformerConfig',
-    'Persformer',
-    'PersformerWrapper',
+    "PersformerConfig",
+    "Persformer",
+    "PersformerWrapper",
     "attention",
-    "pooling_layers"
+    "pooling_layers",
 ]
diff --git a/gdeep/topology_layers/attention/__init__.py b/gdeep/topology_layers/attention/__init__.py
index b0df1c89..d94e669d 100644
--- a/gdeep/topology_layers/attention/__init__.py
+++ b/gdeep/topology_layers/attention/__init__.py
@@ -4,7 +4,7 @@
 
 
 __all__ = [
-    'AttentionBase',
-    'ScaledDotProductAttention',
-    'InducedAttention',
-]
\ No newline at end of file
+    "AttentionBase",
+    "ScaledDotProductAttention",
+    "InducedAttention",
+]
diff --git a/gdeep/topology_layers/attention/attention_base.py b/gdeep/topology_layers/attention/attention_base.py
index 67b9662b..757d82c9 100644
--- a/gdeep/topology_layers/attention/attention_base.py
+++ b/gdeep/topology_layers/attention/attention_base.py
@@ -13,13 +13,13 @@ class AttentionBase(Module, ABC):
     """Base class for attention layers. This class
     can be used in generic transformer models.
     """
+
     def __init__(self, config: PersformerConfig):
         super().__init__()
         self.config = config
-    
+
     @abstractmethod
-    def forward(self,  # type: ignore
-                input: Tensor,
-                attention_mask: Optional[Tensor] = None
-                ) -> Tensor:
+    def forward(
+        self, input: Tensor, attention_mask: Optional[Tensor] = None  # type: ignore
+    ) -> Tensor:
         raise NotImplementedError
diff --git a/gdeep/topology_layers/attention/induced_attention.py b/gdeep/topology_layers/attention/induced_attention.py
index f2bc8876..06b729e6 100644
--- a/gdeep/topology_layers/attention/induced_attention.py
+++ b/gdeep/topology_layers/attention/induced_attention.py
@@ -11,13 +11,12 @@
 
 class InducedAttention(AttentionBase):
     """Class implementing the induced attention"""
+
     def __init__(self, config: PersformerConfig) -> None:
         super().__init__(config)
         raise NotImplementedError
-    
-    def forward(self,  # type: ignore
-                input: Tensor,
-                attention_mask: Optional[Tensor] = None
-                ):
+
+    def forward(
+        self, input: Tensor, attention_mask: Optional[Tensor] = None  # type: ignore
+    ):
         raise NotImplementedError
-    
\ No newline at end of file
diff --git a/gdeep/topology_layers/attention/scaled_dot_product_attention.py b/gdeep/topology_layers/attention/scaled_dot_product_attention.py
index a54102d8..6f5deb9a 100644
--- a/gdeep/topology_layers/attention/scaled_dot_product_attention.py
+++ b/gdeep/topology_layers/attention/scaled_dot_product_attention.py
@@ -14,21 +14,21 @@ class ScaledDotProductAttention(AttentionBase):
     """
     Dot product attention. See https://arxiv.org/abs/1706.03762.
     """
-    def __init__(self,
-                 config: PersformerConfig):
+
+    def __init__(self, config: PersformerConfig):
         super().__init__(config)
-        
-        self.scaled_dot_product_attention = \
-            MultiheadAttention(embed_dim=config.hidden_size,
-                               num_heads=config.num_attention_heads,
-                               dropout=config.attention_probs_dropout_prob,
-                               batch_first=True)
+
+        self.scaled_dot_product_attention = MultiheadAttention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            dropout=config.attention_probs_dropout_prob,
+            batch_first=True,
+        )
         self.dropout = Dropout(config.hidden_dropout_prob)
 
-    def forward(self,  # type: ignore
-                input: Tensor,
-                attention_mask: Optional[Tensor] = None
-                ):
+    def forward(
+        self, input: Tensor, attention_mask: Optional[Tensor] = None  # type: ignore
+    ):
         """
         Forward pass.
         """
@@ -37,6 +37,7 @@ def forward(self,  # type: ignore
         #     attn_mask = torch.concat([attention_mask] * self.config.num_attention_heads, dim=0)
         # else:
         #     attn_mask = None
-        attention_output, _ = self.scaled_dot_product_attention(input, input, input,
-                                                                key_padding_mask=attention_mask)
+        attention_output, _ = self.scaled_dot_product_attention(
+            input, input, input, key_padding_mask=attention_mask
+        )
         return self.dropout(attention_output)
diff --git a/gdeep/topology_layers/attention_factory.py b/gdeep/topology_layers/attention_factory.py
index 93ff7715..b82f93e5 100644
--- a/gdeep/topology_layers/attention_factory.py
+++ b/gdeep/topology_layers/attention_factory.py
@@ -2,36 +2,40 @@
 
 from gdeep.utility.enum_types import AttentionType
 
-from .attention import (AttentionBase, InducedAttention,
-                        ScaledDotProductAttention)
+from .attention import AttentionBase, InducedAttention, ScaledDotProductAttention
 from .persformer_config import PersformerConfig
 
 
-class AttentionFactory():
+class AttentionFactory:
     """
     Factory for creating attention modules.
     """
-    attention_modules: Dict[AttentionType, Callable[[PersformerConfig], AttentionBase]] = {}
-    
+
+    attention_modules: Dict[
+        AttentionType, Callable[[PersformerConfig], AttentionBase]
+    ] = {}
+
     def __init__(self):
-            # Register the attention layers here:
-            self.register_attention__builder(AttentionType.DOT_PRODUCT,
-                                             lambda config: ScaledDotProductAttention(config))
-            self.register_attention__builder(AttentionType.INDUCED_ATTENTION,
-                                                lambda config: InducedAttention(config))
-        
-    
-    def register_attention__builder(self,
-                           attention_type: AttentionType,
-                           attention_module_builder: Callable[[PersformerConfig], AttentionBase]
-                           ) -> None:
+        # Register the attention layers here:
+        self.register_attention__builder(
+            AttentionType.DOT_PRODUCT, lambda config: ScaledDotProductAttention(config)
+        )
+        self.register_attention__builder(
+            AttentionType.INDUCED_ATTENTION, lambda config: InducedAttention(config)
+        )
+
+    def register_attention__builder(
+        self,
+        attention_type: AttentionType,
+        attention_module_builder: Callable[[PersformerConfig], AttentionBase],
+    ) -> None:
         """
         Register an attention module.
         """
         self.attention_modules[attention_type] = attention_module_builder
-        
+
     def build(self, config: PersformerConfig) -> AttentionBase:
-            """
-            Create an attention module.
-            """
-            return self.attention_modules[config.attention_type](config)
+        """
+        Create an attention module.
+        """
+        return self.attention_modules[config.attention_type](config)
diff --git a/gdeep/topology_layers/persformer.py b/gdeep/topology_layers/persformer.py
index 39a79b8d..4dc17add 100644
--- a/gdeep/topology_layers/persformer.py
+++ b/gdeep/topology_layers/persformer.py
@@ -1,4 +1,3 @@
-
 from typing import Optional
 
 import torch
@@ -25,58 +24,62 @@ class Persformer(Module):
         model = Persformer(config)
 
     """
+
     config: PersformerConfig
     embedding_layer: Module
     persformer_blocks: ModuleList
     classifier_layer: Module
     pooling_layer: Module
-    
+
     def __init__(self, config: PersformerConfig):
         super().__init__()
         self.config = config
         self.build_model()
-        
+
     def build_model(self):
         """
         Build the model.
         """
         self.embedding_layer = self._get_embedding_layer()
-        self.persformer_blocks = ModuleList([self._get_persformer_block() 
-                                             for _ in range(self.config.num_attention_layers)])
+        self.persformer_blocks = ModuleList(
+            [
+                self._get_persformer_block()
+                for _ in range(self.config.num_attention_layers)
+            ]
+        )
         self.pooling_layer = self._get_pooling_layer()
         self.classifier_layer = self._get_classifier_layer()
 
     def _get_embedding_layer(self) -> Module:
         return Sequential(
-                            Linear(self.config.input_size, self.config.hidden_size),
-                            get_activation_function(self.config.hidden_act),
-                        )
-        
+            Linear(self.config.input_size, self.config.hidden_size),
+            get_activation_function(self.config.hidden_act),
+        )
+
     def _get_classifier_layer(self) -> Module:
         return Sequential(
-                            Linear(self.config.hidden_size, self.config.hidden_size),
-                            get_activation_function(self.config.hidden_act),
-                            Dropout(self.config.classifier_dropout_prob),
-                            Linear(self.config.hidden_size, self.config.output_size),
-                        )
-        
+            Linear(self.config.hidden_size, self.config.hidden_size),
+            get_activation_function(self.config.hidden_act),
+            Dropout(self.config.classifier_dropout_prob),
+            Linear(self.config.hidden_size, self.config.output_size),
+        )
+
     def _get_persformer_block(self) -> Module:
         return PersformerBlock(self.config)  # type: ignore
-                               
+
     def _get_pooling_layer(self) -> Module:
         return get_pooling_layer(self.config)
 
-    def forward(self,
-                input_batch: Tensor,
-                attention_mask: Optional[Tensor] = None
-                ) -> Tensor:
+    def forward(
+        self, input_batch: Tensor, attention_mask: Optional[Tensor] = None
+    ) -> Tensor:
         """
         Forward pass of the model.
-        
+
         Args:
             input_batch: The input batch. Of shape (batch_size, sequence_length, 2 + num_homology_types)
             attention_mask: The attention mask. Of shape (batch_size, sequence_length)
-        
+
         Returns:
             The logits of the model. Of shape (batch_size, sequence_length, 1)
         """
diff --git a/gdeep/topology_layers/persformer_block.py b/gdeep/topology_layers/persformer_block.py
index f541c64e..6252f67f 100644
--- a/gdeep/topology_layers/persformer_block.py
+++ b/gdeep/topology_layers/persformer_block.py
@@ -10,22 +10,23 @@
 # Type aliases
 Tensor = torch.Tensor
 
+
 class PersformerBlock(Module):
     """
     A persformer block.
     """
-    
+
     config: PersformerConfig
     attention_layer: Module
     feed_forward_layer: Module
     dropout_layer: Module
     layer_norms: Optional[ModuleList]
-    
+
     def __init__(self, config: PersformerConfig):
         super().__init__()
         self.config = config
         self.build_model()
-        
+
     def build_model(self):
         """
         Build the model.
@@ -33,26 +34,33 @@ def build_model(self):
         self.attention_layer = get_attention_layer(self.config)
         if not self.config.use_attention_only:
             self.feed_forward_layer = get_feed_forward_layer(self.config)
-       
-        if self.config.layer_norm_style == LayerNormStyle.PRE_LAYER_NORMALIZATION or \
-            self.config.layer_norm_style == LayerNormStyle.POST_LAYER_NORMALIZATION:
-            self.layer_norms = ModuleList([LayerNorm(self.config.hidden_size, eps=self.config.layer_norm_eps) 
-                                           for _ in range(2)])
-        
-    def forward(self,  # type: ignore
-                input_batch: Tensor,
-                attention_mask: Optional[Tensor] = None
-                ) -> Tensor:
+
+        if (
+            self.config.layer_norm_style == LayerNormStyle.PRE_LAYER_NORMALIZATION
+            or self.config.layer_norm_style == LayerNormStyle.POST_LAYER_NORMALIZATION
+        ):
+            self.layer_norms = ModuleList(
+                [
+                    LayerNorm(self.config.hidden_size, eps=self.config.layer_norm_eps)
+                    for _ in range(2)
+                ]
+            )
+
+    def forward(
+        self,  # type: ignore
+        input_batch: Tensor,
+        attention_mask: Optional[Tensor] = None,
+    ) -> Tensor:
         """
         Forward pass of the model. We implement different layer normalizations in the forward pass,
         see the paper https://arxiv.org/pdf/2002.04745.pdf for details.
-        
+
         Args:
             input_batch:
                 The input batch. Of shape (batch_size, sequence_length, 2 + num_homology_types)
             attention_mask:
                 The attention mask. Of shape (batch_size, sequence_length)
-        
+
         Returns:
             The logits of the model. Of shape (batch_size, sequence_length, 1)
         """
@@ -64,21 +72,19 @@ def forward(self,  # type: ignore
             return self._forward_post_layer_norm(input_batch, attention_mask)
         else:
             raise ValueError(f"Unknown layer norm style {self.config.layer_norm_style}")
-        
-        
-    def _forward_no_layer_norm(self,
-                               input_batch: Tensor,
-                               attention_mask: Optional[Tensor] = None
-                               ) -> Tensor:
+
+    def _forward_no_layer_norm(
+        self, input_batch: Tensor, attention_mask: Optional[Tensor] = None
+    ) -> Tensor:
         """
         Forward pass of the model without layer normalization.
-        
+
         Args:
             input_batch:
                 The input batch. Of shape (batch_size, sequence_length, 2 + num_homology_types)
             attention_mask:
                 The attention mask. Of shape (batch_size, sequence_length)
-        
+
         Returns:
             The logits of the model. Of shape (batch_size, sequence_length, 1)
         """
@@ -86,20 +92,19 @@ def _forward_no_layer_norm(self,
         if not self.config.use_attention_only:
             output = self.feed_forward_layer(output) + output
         return output
-        
-    def _forward_pre_layer_norm(self,
-                                input_batch: Tensor,
-                                attention_mask: Optional[Tensor] = None
-                                ) -> Tensor:
+
+    def _forward_pre_layer_norm(
+        self, input_batch: Tensor, attention_mask: Optional[Tensor] = None
+    ) -> Tensor:
         """
         Forward pass of the model with pre-layer normalization.
-        
+
         Args:
             input_batch:
                 The input batch. Of shape (batch_size, sequence_length, 2 + num_homology_types)
             attention_mask:
                 The attention mask. Of shape (batch_size, sequence_length)
-        
+
         Returns:
             The logits of the model. Of shape (batch_size, sequence_length, 1)
         """
@@ -110,20 +115,19 @@ def _forward_pre_layer_norm(self,
             normalized = self.layer_norms[1](output)
             output = self.feed_forward_layer(normalized) + output
         return output
-    
-    def _forward_post_layer_norm(self,
-                                 input_batch: Tensor,
-                                 attention_mask: Optional[Tensor] = None
-                                 ) -> Tensor:
+
+    def _forward_post_layer_norm(
+        self, input_batch: Tensor, attention_mask: Optional[Tensor] = None
+    ) -> Tensor:
         """
         Forward pass of the model with post-layer normalization.
-        
+
         Args:
             input_batch:
                 The input batch. Of shape (batch_size, sequence_length, 2 + num_homology_types)
             attention_mask:
                 The attention mask. Of shape (batch_size, sequence_length)
-        
+
         Returns:
             The logits of the model. Of shape (batch_size, sequence_length, 1)
         """
diff --git a/gdeep/topology_layers/persformer_config.py b/gdeep/topology_layers/persformer_config.py
index f1da9e39..be75e86c 100644
--- a/gdeep/topology_layers/persformer_config.py
+++ b/gdeep/topology_layers/persformer_config.py
@@ -3,18 +3,31 @@
 from transformers.configuration_utils import PretrainedConfig
 import torch
 import torch.nn as nn
-from torch.nn import Module, MultiheadAttention, Linear, Sequential, LayerNorm, Dropout, ModuleList
+from torch.nn import (
+    Module,
+    MultiheadAttention,
+    Linear,
+    Sequential,
+    LayerNorm,
+    Dropout,
+    ModuleList,
+)
 
-from gdeep.utility.enum_types import ActivationFunction, PoolerType, AttentionType, LayerNormStyle
+from gdeep.utility.enum_types import (
+    ActivationFunction,
+    PoolerType,
+    AttentionType,
+    LayerNormStyle,
+)
 
 # Type aliases
 Tensor = torch.Tensor
-    
+
 
 class PersformerConfig(PretrainedConfig):
     """
     Configuration class to define a persformer model.
-    
+
     Examples::
 
         from gdeep.topological_layers import PersformerConfig, PersformerModel
@@ -26,8 +39,8 @@ class PersformerConfig(PretrainedConfig):
         config = model.config
 
     """
-    
-    input_size: int # input size of the model
+
+    input_size: int  # input size of the model
     output_size: int
     hidden_size: int
     num_attention_layers: int
@@ -44,29 +57,27 @@ class PersformerConfig(PretrainedConfig):
     pooler_type: PoolerType
     use_attention_only: bool
     use_skip_connections_for_persformer_blocks: bool
-         
-    
-    def __init__(self,
-                 input_size: int = 2 + 4,
-                 output_size: int = 2,
-                 hidden_size: int = 32,
-                 num_attention_layers: int = 2,
-                 num_attention_heads: int = 4,
-                 intermediate_size: int = 32,
-                 hidden_act: ActivationFunction = ActivationFunction.GELU,
-                 hidden_dropout_prob: float = 0.1,
-                 attention_probs_dropout_prob: float = 0.1,
-                 layer_norm_eps: float = 1e-12,
-                 classifier_dropout_prob: float = 0.1,
-                 use_layer_norm: LayerNormStyle = \
-                     LayerNormStyle.NO_LAYER_NORMALIZATION,
-                 attention_type: AttentionType = \
-                     AttentionType.DOT_PRODUCT,
-                 pooler_type: PoolerType = PoolerType.ATTENTION,
-                 use_attention_only: bool = False,
-                 use_skip_connections_for_persformer_blocks=False,
-                 **kwargs
-                 ):
+
+    def __init__(
+        self,
+        input_size: int = 2 + 4,
+        output_size: int = 2,
+        hidden_size: int = 32,
+        num_attention_layers: int = 2,
+        num_attention_heads: int = 4,
+        intermediate_size: int = 32,
+        hidden_act: ActivationFunction = ActivationFunction.GELU,
+        hidden_dropout_prob: float = 0.1,
+        attention_probs_dropout_prob: float = 0.1,
+        layer_norm_eps: float = 1e-12,
+        classifier_dropout_prob: float = 0.1,
+        use_layer_norm: LayerNormStyle = LayerNormStyle.NO_LAYER_NORMALIZATION,
+        attention_type: AttentionType = AttentionType.DOT_PRODUCT,
+        pooler_type: PoolerType = PoolerType.ATTENTION,
+        use_attention_only: bool = False,
+        use_skip_connections_for_persformer_blocks=False,
+        **kwargs
+    ):
         super().__init__(**kwargs)  # type: ignore
         self.input_size = input_size
         self.output_size = output_size
@@ -83,6 +94,6 @@ def __init__(self,
         self.attention_type = attention_type
         self.pooler_type = pooler_type
         self.use_attention_only = use_attention_only
-        self.use_skip_connections_for_persformer_blocks = use_skip_connections_for_persformer_blocks
-
-
+        self.use_skip_connections_for_persformer_blocks = (
+            use_skip_connections_for_persformer_blocks
+        )
diff --git a/gdeep/topology_layers/persformer_wrapper.py b/gdeep/topology_layers/persformer_wrapper.py
index 1dd71493..c2de4ed3 100644
--- a/gdeep/topology_layers/persformer_wrapper.py
+++ b/gdeep/topology_layers/persformer_wrapper.py
@@ -1,9 +1,12 @@
-
 from typing import Optional
 
 import torch
-from gdeep.utility.enum_types import (ActivationFunction, AttentionType,
-                                      LayerNormStyle, PoolerType)
+from gdeep.utility.enum_types import (
+    ActivationFunction,
+    AttentionType,
+    LayerNormStyle,
+    PoolerType,
+)
 from torch.nn import Module
 
 from .persformer import Persformer
@@ -16,19 +19,14 @@ class PersformerWrapper(Module):
     """The wrapper for persformer to allow compatibility
     with the HPO classes.
     """
+
     config: PersformerConfig
     model: Module
-    
-    def __init__(self,
-                 **kwargs):
+
+    def __init__(self, **kwargs):
         super().__init__()
-        self.config = PersformerConfig(
-            **kwargs
-        )
+        self.config = PersformerConfig(**kwargs)
         self.model = Persformer(self.config)
-    
-    def forward(self, 
-                input: Tensor,
-                attention_mask: Optional[Tensor] = None
-                ):
+
+    def forward(self, input: Tensor, attention_mask: Optional[Tensor] = None):
         return self.model(input, attention_mask)
diff --git a/gdeep/topology_layers/pooling_layers/__init__.py b/gdeep/topology_layers/pooling_layers/__init__.py
index 5b19f240..35d5d6ec 100644
--- a/gdeep/topology_layers/pooling_layers/__init__.py
+++ b/gdeep/topology_layers/pooling_layers/__init__.py
@@ -8,4 +8,4 @@
     "MaxPoolingLayer",
     "MeanPoolingLayer",
     "SumPoolingLayer",
-]
\ No newline at end of file
+]
diff --git a/gdeep/topology_layers/pooling_layers/attention_pooling_layer.py b/gdeep/topology_layers/pooling_layers/attention_pooling_layer.py
index e501ba96..74baf20d 100644
--- a/gdeep/topology_layers/pooling_layers/attention_pooling_layer.py
+++ b/gdeep/topology_layers/pooling_layers/attention_pooling_layer.py
@@ -14,38 +14,39 @@ class AttentionPoolingLayer(Module):
     """This class implements the attention mechanism
     with a pooling layer to enforce permutation
     invariance"""
+
     config: PersformerConfig
-    
+
     def __init__(self, config: PersformerConfig):
         super().__init__()
         self.config = config
-        self.query = nn.parameter.Parameter(torch.Tensor(1, config.hidden_size),
-                                              requires_grad=True)
-        self.scaled_dot_product_attention = \
-            MultiheadAttention(embed_dim=config.hidden_size,
-                               num_heads=config.num_attention_heads,
-                               dropout=config.attention_probs_dropout_prob,
-                               batch_first=True)
-
-        
-    def forward(self,
-                input_batch: Tensor,
-                attention_mask: Optional[Tensor] = None
-                ) -> Tensor:
+        self.query = nn.parameter.Parameter(
+            torch.Tensor(1, config.hidden_size), requires_grad=True
+        )
+        self.scaled_dot_product_attention = MultiheadAttention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            dropout=config.attention_probs_dropout_prob,
+            batch_first=True,
+        )
+
+    def forward(
+        self, input_batch: Tensor, attention_mask: Optional[Tensor] = None
+    ) -> Tensor:
         """
         Forward pass of the model.
-        
+
         Args:
             input_batch:
                 The input batch. Of shape (batch_size, sequence_length, hidden_size)
-            
+
         Returns:
             The pooled output. Of shape (batch_size, hidden_size)
         """
         output, _ = self.scaled_dot_product_attention(
-                                                 self.query.expand(input_batch.shape[0], -1, -1),
-                                                 input_batch,
-                                                 input_batch,
-                                                 key_padding_mask=attention_mask
-                                                 )
+            self.query.expand(input_batch.shape[0], -1, -1),
+            input_batch,
+            input_batch,
+            key_padding_mask=attention_mask,
+        )
         return output.squeeze(dim=1)
diff --git a/gdeep/topology_layers/pooling_layers/max_pooling_layer.py b/gdeep/topology_layers/pooling_layers/max_pooling_layer.py
index a9b083d4..72ad6cf1 100644
--- a/gdeep/topology_layers/pooling_layers/max_pooling_layer.py
+++ b/gdeep/topology_layers/pooling_layers/max_pooling_layer.py
@@ -11,23 +11,23 @@
 
 class MaxPoolingLayer(Module):
     """Implementation of the max pooling layer"""
+
     config: PersformerConfig
-    
+
     def __init__(self, config: PersformerConfig):
         super().__init__()
         self.config = config
-        
-    def forward(self,
-                input_batch: Tensor,
-                attention_mask: Optional[Tensor] = None
-                ) -> Tensor:
+
+    def forward(
+        self, input_batch: Tensor, attention_mask: Optional[Tensor] = None
+    ) -> Tensor:
         """
         Forward pass of the model.
-        
+
         Args:
             input_batch:
                 The input batch. Of shape (batch_size, sequence_length, hidden_size)
-            
+
         Returns:
             The logits of the model. Of shape (batch_size, sequence_length, 1)
         """
@@ -35,7 +35,7 @@ def forward(self,
         if attention_mask is not None:
             output = input_batch * attention_mask.unsqueeze(2)
         else:
-            output = input_batch 
+            output = input_batch
         # Apply the max pooling layer
         output = output.max(dim=-2)[0]
         return output
diff --git a/gdeep/topology_layers/pooling_layers/mean_pooling_layer.py b/gdeep/topology_layers/pooling_layers/mean_pooling_layer.py
index 19534f51..f662dbd2 100644
--- a/gdeep/topology_layers/pooling_layers/mean_pooling_layer.py
+++ b/gdeep/topology_layers/pooling_layers/mean_pooling_layer.py
@@ -8,33 +8,33 @@
 # Type aliases
 Tensor = torch.Tensor
 
+
 class MeanPoolingLayer(Module):
-        
-        config: PersformerConfig
-        
-        def __init__(self, config: PersformerConfig):
-            super().__init__()
-            self.config = config
-            
-        def forward(self,
-                    input_batch: Tensor,
-                    attention_mask: Optional[Tensor] = None
-                    ) -> Tensor:
-            """
-            Forward pass of the model.
-            
-            Args:
-                input_batch:
-                    The input batch. Of shape (batch_size, sequence_length, hidden_size)
-                
-            Returns:
-                The pooled output. Of shape (batch_size, hidden_size)
-            """
-            # Initialize the output tensor
-            if attention_mask is not None:
-                output = input_batch * attention_mask.unsqueeze(dim=-1)
-                output = output.sum(dim=-2) / attention_mask.sum(dim=-1, keepdim=True)
-                return output
-            else:
-                output = input_batch
-                return output.mean(dim=-2)
+
+    config: PersformerConfig
+
+    def __init__(self, config: PersformerConfig):
+        super().__init__()
+        self.config = config
+
+    def forward(
+        self, input_batch: Tensor, attention_mask: Optional[Tensor] = None
+    ) -> Tensor:
+        """
+        Forward pass of the model.
+
+        Args:
+            input_batch:
+                The input batch. Of shape (batch_size, sequence_length, hidden_size)
+
+        Returns:
+            The pooled output. Of shape (batch_size, hidden_size)
+        """
+        # Initialize the output tensor
+        if attention_mask is not None:
+            output = input_batch * attention_mask.unsqueeze(dim=-1)
+            output = output.sum(dim=-2) / attention_mask.sum(dim=-1, keepdim=True)
+            return output
+        else:
+            output = input_batch
+            return output.mean(dim=-2)
diff --git a/gdeep/topology_layers/pooling_layers/sum_pooling_layer.py b/gdeep/topology_layers/pooling_layers/sum_pooling_layer.py
index 6f2457ad..ce38f647 100644
--- a/gdeep/topology_layers/pooling_layers/sum_pooling_layer.py
+++ b/gdeep/topology_layers/pooling_layers/sum_pooling_layer.py
@@ -10,32 +10,31 @@
 
 
 class SumPoolingLayer(Module):
-        
-        config: PersformerConfig
-        
-        def __init__(self, config: PersformerConfig):
-            super().__init__()
-            self.config = config
-            
-        def forward(self,
-                    input_batch: Tensor,
-                    attention_mask: Optional[Tensor] = None
-                    ) -> Tensor:
-            """
-            Forward pass of the model.
-            
-            Args:
-                input_batch:
-                    The input batch. Of shape (batch_size, sequence_length, hidden_size)
-                
-            Returns:
-                The pooled output. Of shape (batch_size, hidden_size)
-            """
-            # Initialize the output tensor
-            if attention_mask is not None:
-                output = input_batch * attention_mask.unsqueeze(dim=-1)
-            else:
-                output = input_batch
-            # Apply the max pooling layer
-            output = output.sum(dim=-2)
-            return output
+
+    config: PersformerConfig
+
+    def __init__(self, config: PersformerConfig):
+        super().__init__()
+        self.config = config
+
+    def forward(
+        self, input_batch: Tensor, attention_mask: Optional[Tensor] = None
+    ) -> Tensor:
+        """
+        Forward pass of the model.
+
+        Args:
+            input_batch:
+                The input batch. Of shape (batch_size, sequence_length, hidden_size)
+
+        Returns:
+            The pooled output. Of shape (batch_size, hidden_size)
+        """
+        # Initialize the output tensor
+        if attention_mask is not None:
+            output = input_batch * attention_mask.unsqueeze(dim=-1)
+        else:
+            output = input_batch
+        # Apply the max pooling layer
+        output = output.sum(dim=-2)
+        return output
diff --git a/gdeep/topology_layers/utility.py b/gdeep/topology_layers/utility.py
index 6de32110..35d840e6 100644
--- a/gdeep/topology_layers/utility.py
+++ b/gdeep/topology_layers/utility.py
@@ -1,11 +1,12 @@
-
 import torch.nn as nn
 from torch.nn import Dropout, Linear, Module, Sequential
 
-from gdeep.topology_layers.pooling_layers import (AttentionPoolingLayer,
-                                                  MaxPoolingLayer,
-                                                  MeanPoolingLayer,
-                                                  SumPoolingLayer)
+from gdeep.topology_layers.pooling_layers import (
+    AttentionPoolingLayer,
+    MaxPoolingLayer,
+    MeanPoolingLayer,
+    SumPoolingLayer,
+)
 from gdeep.utility.enum_types import ActivationFunction, PoolerType
 from .attention_factory import AttentionFactory
 from .persformer_config import PersformerConfig
@@ -14,10 +15,10 @@
 def get_pooling_layer(config: PersformerConfig) -> Module:
     """
     Get the pooling layer.
-    
+
     Args:
         config: The configuration of the model.
-        
+
     Returns:
         The pooling layer.
     """
@@ -39,26 +40,17 @@ def get_feed_forward_layer(config: PersformerConfig) -> Module:
     """
     feed_forward_layer = Sequential()
     feed_forward_layer.add_module(
-                                    "intermediate",
-                                    Linear(config.hidden_size, config.intermediate_size)
-                                )
+        "intermediate", Linear(config.hidden_size, config.intermediate_size)
+    )
     feed_forward_layer.add_module(
-                                    "activation",
-                                    get_activation_function(config.hidden_act)
-                                )
-    feed_forward_layer.add_module(
-                                    "dropout",
-                                    Dropout(config.hidden_dropout_prob)
-                                )
-    
-    feed_forward_layer.add_module(
-                                    "output",
-                                    Linear(config.intermediate_size, config.hidden_size)
-                                    )
+        "activation", get_activation_function(config.hidden_act)
+    )
+    feed_forward_layer.add_module("dropout", Dropout(config.hidden_dropout_prob))
+
     feed_forward_layer.add_module(
-                                    "dropout",
-                                    Dropout(config.hidden_dropout_prob)
-                                )
+        "output", Linear(config.intermediate_size, config.hidden_size)
+    )
+    feed_forward_layer.add_module("dropout", Dropout(config.hidden_dropout_prob))
     return feed_forward_layer
 
 
diff --git a/gdeep/trainer/metrics.py b/gdeep/trainer/metrics.py
index c627b735..a2677428 100644
--- a/gdeep/trainer/metrics.py
+++ b/gdeep/trainer/metrics.py
@@ -18,7 +18,7 @@ def accuracy(prediction: Tensor, y: Tensor) -> float:
     Returns:
         float:
             the value of the accuracy
-        """
+    """
     correct: float = 0.0
     try:
         correct += (prediction.argmax(1) == y).to(torch.float).sum().item()
diff --git a/gdeep/trainer/trainer.py b/gdeep/trainer/trainer.py
index f2f6eba7..1a5443dc 100644
--- a/gdeep/trainer/trainer.py
+++ b/gdeep/trainer/trainer.py
@@ -125,9 +125,12 @@ def forward(self, x):
         pipe.train(SGD, 2, True, {"lr": 0.001}, n_accumulated_grads=5)
 
     """
+
     scheduler: Optional[_LRScheduler]
     writer: Optional[SummaryWriter]
-    registered_hook: Optional[Callable[[int, Optimizer, ModelExtractor, Optional[SummaryWriter]], Any]] = None
+    registered_hook: Optional[
+        Callable[[int, Optimizer, ModelExtractor, Optional[SummaryWriter]], Any]
+    ] = None
 
     def __init__(
         self,
@@ -145,7 +148,7 @@ def __init__(
         self.train_epoch = 0
         self.val_epoch = 0
         self.best_val_loss = np.inf
-        self.best_val_acc = 0.
+        self.best_val_acc = 0.0
         self.loss_fn = loss_fn
         if training_metric:
             self.training_metric = training_metric
@@ -240,8 +243,9 @@ def _optimisation_step(
             )
         return epoch_loss
 
-    def _send_to_device(self, x: Union[Tensor, List[Tensor]], y: Tensor) \
-            -> Tuple[Tensor, Union[Tensor, List[Tensor]], Tensor]:
+    def _send_to_device(
+        self, x: Union[Tensor, List[Tensor]], y: Tensor
+    ) -> Tuple[Tensor, Union[Tensor, List[Tensor]], Tensor]:
         """use this private method to send the
         ``x`` and ``y`` to the ``DEVICE``.
 
@@ -289,6 +293,7 @@ def closure() -> Tensor:
                 loss2 = self.loss_fn(self.model(X), y)
                 loss2.backward()
                 return loss2
+
             pred, X, y = self._send_to_device(X, y)
             batch_metric = self.training_metric(pred, y)
             metric_list.append(batch_metric)
@@ -334,7 +339,9 @@ def closure() -> Tensor:
         return epoch_metric, epoch_loss
 
     def _train_loop(
-        self, dl_tr: DataLoader[Tuple[Union[Tensor, List[Tensor]], Tensor]], writer_tag: str = ""
+        self,
+        dl_tr: DataLoader[Tuple[Union[Tensor, List[Tensor]], Tensor]],
+        writer_tag: str = "",
     ) -> Tuple[float, float]:
         """private method to run a single training
         loop
@@ -369,7 +376,9 @@ def _train_loop(
         return epoch_metric, epoch_loss
 
     def _val_loop(
-        self, dl_val: DataLoader[Tuple[Union[Tensor, List[Tensor]], Tensor]], writer_tag: str = ""
+        self,
+        dl_val: DataLoader[Tuple[Union[Tensor, List[Tensor]], Tensor]],
+        writer_tag: str = "",
     ) -> Tuple[float, float]:
         """private method to run a single validation
         loop
@@ -531,7 +540,7 @@ def _init_optimizer_and_scheduler(
         lr_scheduler: Optional[Type[_LRScheduler]] = None,
         scheduler_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        """Reset or maintain the LR scheduler and the 
+        """Reset or maintain the LR scheduler and the
         optimizer depending on the training"""
         if not (self.check_has_trained and keep_training):
             # reset the model weights
@@ -890,7 +899,7 @@ def _training_loops(
         writer_tag: str = "",
     ) -> Tuple[float, float]:
         """private method to run the trainign loops
-        
+
         Args:
             n_epochs:
                 number of training epochs
@@ -1103,9 +1112,13 @@ def map_fun_custom(index, flags):
                     # Compute prediction and loss
                     pred = model2(X)
                     try:
-                        correct += (pred.argmax(1) == y).to(torch.float).sum().item()  # noqa
+                        correct += (
+                            (pred.argmax(1) == y).to(torch.float).sum().item()
+                        )  # noqa
                     except RuntimeError:
-                        correct += (pred.argmax(2) == y).to(torch.float).sum().item()  # noqa
+                        correct += (
+                            (pred.argmax(2) == y).to(torch.float).sum().item()
+                        )  # noqa
                     loss = self.loss_fn(pred, y)
                     # Save to tensorboard
                     # self.writer.add_scalar("Parallel" + "/Loss/train",
@@ -1183,17 +1196,19 @@ def map_fun_custom(index, flags):
         return self.val_loss, self.val_acc
 
     def evaluate_classification(
-        self, num_class: int = None, dl: DataLoader[Tuple[Union[Tensor, List[Tensor]], Tensor]] = None
+        self,
+        num_class: int = None,
+        dl: DataLoader[Tuple[Union[Tensor, List[Tensor]], Tensor]] = None,
     ) -> Tuple[float, float, np.ndarray]:
         """Method to evaluate the performance of the model.
-        
+
         Args:
             num_class:
                 number of classes
             dl :
                 the Dataloader to evaluate. If ``None``,
                 we use the training dataloader in ``self``
-                
+
         Returns:
             (float, float, 2darray):
                 the accuracy, loss and confusion matrix.
@@ -1223,7 +1238,10 @@ def evaluate_classification(
         return epoch_metric, epoch_loss, confusion_matrix
 
     def register_pipe_hook(
-        self, callable: Callable[[int, Optimizer, ModelExtractor, Optional[SummaryWriter]], Any]
+        self,
+        callable: Callable[
+            [int, Optimizer, ModelExtractor, Optional[SummaryWriter]], Any
+        ],
     ) -> None:
         """This method registers a function that
         will be called after each trainign step.
@@ -1240,7 +1258,11 @@ def register_pipe_hook(
         self.registered_hook = callable
 
     def _run_pipe_hook(
-        self, epoch: int, optim: Optimizer, me: ModelExtractor, writer: Optional[SummaryWriter] = None
+        self,
+        epoch: int,
+        optim: Optimizer,
+        me: ModelExtractor,
+        writer: Optional[SummaryWriter] = None,
     ) -> None:
         """private method that runs the hooked
         function at every epoch, after the single training loop"""
diff --git a/gdeep/utility/_typing_utils.py b/gdeep/utility/_typing_utils.py
index 9ece1c86..d67fbf4a 100644
--- a/gdeep/utility/_typing_utils.py
+++ b/gdeep/utility/_typing_utils.py
@@ -10,12 +10,12 @@
 def torch_transform(
     transform: Union[Callable[[Tensor], Tensor], Callable[[ndarray], ndarray]]
 ) -> Callable[[Tensor], Tensor]:
-    """ Transforms a numpy array transform to a torch transform. 
-    
+    """Transforms a numpy array transform to a torch transform.
+
     Args:
         transform: Either a callable that takes a numpy array and returns a
             numpy array or a callable that takes a tensor and returns a tensor.
-        
+
     Returns:
         The torch transform.
     """
@@ -38,8 +38,8 @@ def torch_transform(
 
 
 def get_parameter_types(func: Callable) -> List[Any]:
-    """ Returns a list of the types of the parameters of a function. 
-    
+    """Returns a list of the types of the parameters of a function.
+
     Args:
         func: The function to get the types of the parameters of.
     Returns:
@@ -49,11 +49,11 @@ def get_parameter_types(func: Callable) -> List[Any]:
 
 
 def get_return_type(func: Callable) -> Any:
-    """ Returns the type of the return value of a function.
-     
+    """Returns the type of the return value of a function.
+
     Args:
         func: The function to get the type of the return value of.
-        
+
     Returns:
         The type of the return value of the function.
     """
diff --git a/gdeep/utility/enum_types.py b/gdeep/utility/enum_types.py
index ab6f24d9..9cb4f5e0 100644
--- a/gdeep/utility/enum_types.py
+++ b/gdeep/utility/enum_types.py
@@ -1,35 +1,41 @@
 from enum import Enum, auto
 
+
 class PoolerType(Enum):
     ATTENTION = auto()
     MAX = auto()
     MEAN = auto()
     SUM = auto()
 
+
 class LayerNormStyle(Enum):
     """
     The style of layer normalization.
     """
+
     NO_LAYER_NORMALIZATION = auto()
     PRE_LAYER_NORMALIZATION = auto()
     POST_LAYER_NORMALIZATION = auto()
-    
+
+
 class AttentionType(Enum):
     """
     The type of attention.
     """
+
     NO_ATTENTION = auto()
     DOT_PRODUCT = auto()
     UNNORMALIZED_DOT_PRODUCT = auto()
     INDUCED_ATTENTION = auto()
     FOURIER_MIXER = auto()
 
+
 class ActivationFunction(Enum):
     """
     The activation function.
     """
+
     RELU = auto()
     GELU = auto()
     SELU = auto()
     MISH = auto()
-    
\ No newline at end of file
diff --git a/gdeep/utility/extended_persistence/__init__.py b/gdeep/utility/extended_persistence/__init__.py
index a6bb77a5..5fe0ae54 100644
--- a/gdeep/utility/extended_persistence/__init__.py
+++ b/gdeep/utility/extended_persistence/__init__.py
@@ -1,5 +1,4 @@
 from .utils import graph_extended_persistence_hks, graph_extended_persistence_gudhi
 
 
-__all__ = ['graph_extended_persistence_hks',
-           'graph_extended_persistence_gudhi']
+__all__ = ["graph_extended_persistence_hks", "graph_extended_persistence_gudhi"]
diff --git a/gdeep/utility/extended_persistence/utils.py b/gdeep/utility/extended_persistence/utils.py
index 9ccc058e..1693312f 100644
--- a/gdeep/utility/extended_persistence/utils.py
+++ b/gdeep/utility/extended_persistence/utils.py
@@ -15,34 +15,34 @@
 
 
 # Compute the heat kernel signature of a graph.
-def _heat_kernel_signature(adj_mat: Array,
-                           diffusion_parameter: float = 1.0) -> Array:
+def _heat_kernel_signature(adj_mat: Array, diffusion_parameter: float = 1.0) -> Array:
     """Given a graph adjacency matrix, compute the heat kernel signature.
-    
+
     Args:
         adj_mat:
             The adjacency matrix of the graph.
         diffusion_parameter:
             The diffusion parameter of the heat kernel.
-            
+
     Returns:
         ndarray:
             The heat kernel signature of the graph. The shape of the array is
             (num_vertices).
     """
     eigenvals, eigenvectors = _get_eigenvalues_eigenvectors(adj_mat)
-    hks = (np.square(eigenvectors) * np.exp(-diffusion_parameter * eigenvals))\
-        .sum(axis=1)
+    hks = (np.square(eigenvectors) * np.exp(-diffusion_parameter * eigenvals)).sum(
+        axis=1
+    )
     return hks
 
 
 def _get_eigenvalues_eigenvectors(adj_mat) -> Tuple[Array, Array]:
     """Compute the eigenvector and eigenvalue of the adjacency matrix.
-    
+
     Args:
         adj_mat:
             The adjacency matrix of the graph.
-        
+
     Returns:
         (ndarray, ndarray):
             The eigenvectors and eigenvalues of the adjacency matrix.
@@ -54,17 +54,17 @@ def _get_eigenvalues_eigenvectors(adj_mat) -> Tuple[Array, Array]:
     return eigenvalues, eigenvectors
 
 
-def graph_extended_persistence_hks(adj_mat: Array,
-                                   diffusion_parameter: float = 1.0) \
-                                       -> OneHotEncodedPersistenceDiagram:
+def graph_extended_persistence_hks(
+    adj_mat: Array, diffusion_parameter: float = 1.0
+) -> OneHotEncodedPersistenceDiagram:
     """Compute the extended persistence of a graph.
-    
+
     Args:
         adj_mat:
             The adjacency matrix of the graph.
         diffusion_parameter:
             The diffusion parameter of the heat kernel.
-            
+
     Returns:
         ndarray:
             The extended persistence of the graph.
@@ -78,31 +78,71 @@ def graph_extended_persistence_hks(adj_mat: Array,
     )
 
 
-def graph_extended_persistence_gudhi(A: Array,
-                                     filtration_val: Array)\
-    -> Tuple[Array, Array, Array, Array]:
-    
+def graph_extended_persistence_gudhi(
+    A: Array, filtration_val: Array
+) -> Tuple[Array, Array, Array, Array]:
+
     # TODO: Rewrite this function
     num_vertices = A.shape[0]
     (xs, ys) = np.where(np.triu(A))
     st = gd.SimplexTree()  # type: ignore
     for i in range(num_vertices):
         st.insert([i], filtration=-1e10)
-    for idx, x in enumerate(xs):        
+    for idx, x in enumerate(xs):
         st.insert([x, ys[idx]], filtration=-1e10)
     for i in range(num_vertices):
         st.assign_filtration([i], filtration_val[i])
     st.make_filtration_non_decreasing()
     st.extend_filtration()
     extended_persistence = st.extended_persistence()
-    dgmOrd0, dgmRel1, dgmExt0, dgmExt1 = extended_persistence[0], extended_persistence[1], \
-        extended_persistence[2], extended_persistence[3]
-    dgmOrd0 = np.vstack([np.array([[ min(p[1][0],p[1][1]), max(p[1][0],p[1][1]) ]]) \
-        for p in dgmOrd0 if p[0] == 0]) if len(dgmOrd0) else np.empty([0,2])
-    dgmRel1 = np.vstack([np.array([[ min(p[1][0],p[1][1]), max(p[1][0],p[1][1]) ]]) \
-        for p in dgmRel1 if p[0] == 1]) if len(dgmRel1) else np.empty([0,2])
-    dgmExt0 = np.vstack([np.array([[ min(p[1][0],p[1][1]), max(p[1][0],p[1][1]) ]]) \
-        for p in dgmExt0 if p[0] == 0]) if len(dgmExt0) else np.empty([0,2])
-    dgmExt1 = np.vstack([np.array([[ min(p[1][0],p[1][1]), max(p[1][0],p[1][1]) ]]) \
-        for p in dgmExt1 if p[0] == 1]) if len(dgmExt1) else np.empty([0,2])
+    dgmOrd0, dgmRel1, dgmExt0, dgmExt1 = (
+        extended_persistence[0],
+        extended_persistence[1],
+        extended_persistence[2],
+        extended_persistence[3],
+    )
+    dgmOrd0 = (
+        np.vstack(
+            [
+                np.array([[min(p[1][0], p[1][1]), max(p[1][0], p[1][1])]])
+                for p in dgmOrd0
+                if p[0] == 0
+            ]
+        )
+        if len(dgmOrd0)
+        else np.empty([0, 2])
+    )
+    dgmRel1 = (
+        np.vstack(
+            [
+                np.array([[min(p[1][0], p[1][1]), max(p[1][0], p[1][1])]])
+                for p in dgmRel1
+                if p[0] == 1
+            ]
+        )
+        if len(dgmRel1)
+        else np.empty([0, 2])
+    )
+    dgmExt0 = (
+        np.vstack(
+            [
+                np.array([[min(p[1][0], p[1][1]), max(p[1][0], p[1][1])]])
+                for p in dgmExt0
+                if p[0] == 0
+            ]
+        )
+        if len(dgmExt0)
+        else np.empty([0, 2])
+    )
+    dgmExt1 = (
+        np.vstack(
+            [
+                np.array([[min(p[1][0], p[1][1]), max(p[1][0], p[1][1])]])
+                for p in dgmExt1
+                if p[0] == 1
+            ]
+        )
+        if len(dgmExt1)
+        else np.empty([0, 2])
+    )
     return dgmOrd0, dgmExt0, dgmRel1, dgmExt1
diff --git a/gdeep/utility/optimisation/__init__.py b/gdeep/utility/optimisation/__init__.py
index f59c3e13..106aa441 100644
--- a/gdeep/utility/optimisation/__init__.py
+++ b/gdeep/utility/optimisation/__init__.py
@@ -2,9 +2,4 @@
 from .sam import SAM, MissingClosureError, SAMOptimizer
 
 
-__all__ = [
-    'PersistenceGradient',
-    'SAM',
-    'SAMOptimizer',
-    'MissingClosureError'
-    ]
+__all__ = ["PersistenceGradient", "SAM", "SAMOptimizer", "MissingClosureError"]
diff --git a/gdeep/utility/optimisation/persistence_grad.py b/gdeep/utility/optimisation/persistence_grad.py
index ef1da40b..1b6ef9c2 100644
--- a/gdeep/utility/optimisation/persistence_grad.py
+++ b/gdeep/utility/optimisation/persistence_grad.py
@@ -69,16 +69,19 @@ class PersistenceGradient:
         pg.sgd(X, n_epochs=4, lr=0.4)
 
     """
+
     dist_mat: Tensor
 
-    def __init__(self,
-                 homology_dimensions: Optional[List[int]],
-                 zeta: float = 0.5,
-                 collapse_edges: bool = False,
-                 max_edge_length: float = np.inf,
-                 approx_digits: int = 6,
-                 metric: str = "euclidean",
-                 directed: bool = False) -> None:
+    def __init__(
+        self,
+        homology_dimensions: Optional[List[int]],
+        zeta: float = 0.5,
+        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
@@ -92,20 +95,16 @@ def __init__(self,
             self.homology_dimensions = [0, 1]
 
     @staticmethod
-    def powerset(iterable: Iterator,
-                 max_length: int) -> Iterator:
+    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))
+        return chain.from_iterable(combinations(s, r) for r in range(0, max_length + 1))
 
     @staticmethod
-    def _parallel_apply_along_axis(func1d: Callable,
-                                   axis: int,
-                                   arr: Array,
-                                   *args: Any,
-                                   **kwargs: Any) -> Array:
+    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
         cores.
@@ -116,10 +115,12 @@ def _parallel_apply_along_axis(func1d: Callable,
         effective_axis = 1 if axis == 0 else axis
         if effective_axis != axis:
             arr = arr.swapaxes(axis, effective_axis)
-        n_processes = min(len(arr)//2, multiprocessing.cpu_count())
+        n_processes = min(len(arr) // 2, multiprocessing.cpu_count())
         # Chunks for the mapping (only a few chunks):
-        chunks = [(func1d, effective_axis, sub_arr, args, kwargs)
-                  for sub_arr in np.array_split(arr, n_processes)]
+        chunks = [
+            (func1d, effective_axis, sub_arr, args, kwargs)
+            for sub_arr in np.array_split(arr, n_processes)
+        ]
         pool = multiprocessing.Pool(n_processes)
         individual_results = pool.map(unpacking_apply_along_axis, chunks)  # type: ignore
         # Freeing the workers:
@@ -131,20 +132,23 @@ 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))),  # type: ignore
-                         max(self.homology_dimensions) + 2))[1:]
+        simplices = list(
+            self.powerset(
+                list(range(0, len(x))),  # type: ignore
+                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)
         # the current computation bottleneck
         if len(simplices_array) > 10000000:
             pairs_of_indices = self._parallel_apply_along_axis(
-                                    _combinations_with_single, 1,
-                                    simplices_array, comb_number)
+                _combinations_with_single, 1, simplices_array, comb_number
+            )
         else:
-            pairs_of_indices = np.apply_along_axis(_combinations_with_single,
-                                                   1,
-                                                   simplices_array,
-                                                   comb_number)
+            pairs_of_indices = np.apply_along_axis(
+                _combinations_with_single, 1, simplices_array, comb_number
+            )
         return torch.tensor(pairs_of_indices, dtype=torch.int64)
 
     def phi(self, x: Tensor) -> Tensor:
@@ -174,10 +178,14 @@ def phi(self, x: Tensor) -> Tensor:
         # lista = [max([dist_mat[pair] for pair in pairs]) for pairs in
         #    simplicial_pairs if max([dist_mat[pair] for pair in pairs])
         #    <= self.max_edge_length]
-        lista = torch.max((torch.gather(torch.index_select(self.dist_mat,
-                                                           0, ks),
-                          1, js.reshape(-1, 1))).reshape(-1,
-                                                         comb_number), 1)[0]
+        lista = torch.max(
+            (
+                torch.gather(
+                    torch.index_select(self.dist_mat, 0, ks), 1, js.reshape(-1, 1)
+                )
+            ).reshape(-1, comb_number),
+            1,
+        )[0]
         lista = torch.sort(lista)[0]
         return lista
 
@@ -232,27 +240,27 @@ def _persistence(self, x: Tensor) -> Tuple[Array, Tensor, Tensor]:
         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: Tensor = torch.round(phi*num_approx)*inv_num_approx
+        inv_num_approx = 10 ** (-self.approx_digits)
+        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_ls: List = []
         for i in approx_pairs:
-            indices_in_phi_of_pairs_ls += \
+            indices_in_phi_of_pairs_ls += (
                 torch.nonzero(phi_approx == i).flatten().detach().tolist()
+            )
         indices_in_phi_of_pairs: Set = set(indices_in_phi_of_pairs_ls)
-        persistence_pairs_array = -np.ones((len(approx_pairs), 2),
-                                           dtype=np.int32)
+        persistence_pairs_array = -np.ones((len(approx_pairs), 2), dtype=np.int32)
         for i in indices_in_phi_of_pairs:
             try:
-                temp_index: int = approx_pairs.index(round(phi[i].item(),  # type: ignore
-                                                     self.approx_digits))
-                approx_pairs[temp_index] = float('inf')
-                persistence_pairs_array[temp_index//2,
-                                        temp_index % 2] = int(i)
+                temp_index: int = approx_pairs.index(
+                    round(phi[i].item(), self.approx_digits)  # type: ignore
+                )
+                approx_pairs[temp_index] = float("inf")
+                persistence_pairs_array[temp_index // 2, temp_index % 2] = int(i)
             except ValueError:
                 pass
         return persistence_pairs_array, phi, homology_dims
-    
+
     def _computing_persistence_with_gph(self, xx: Tensor) -> List:
         """This method accepts the pointcloud and returns the
         persistence diagram in the following form
@@ -270,28 +278,36 @@ def _computing_persistence_with_gph(self, xx: Tensor) -> List:
         Returns:
             list of shape (n, 3):
                 Persistence pairs (corresponding to the
-                first 2 dimensions) where the last dimension 
+                first 2 dimensions) where the last dimension
                 contains the homology dimension
         """
-        output = ripser_parallel(xx.detach().numpy(),
-                                 maxdim=max(self.homology_dimensions),
-                                 thresh=self.max_edge_length,
-                                 coeff=2,
-                                 metric=self.metric,
-                                 collapse_edges=self.collapse_edges,
-                                 n_threads=-1,
-                                 return_generators=True)
-        
+        output = ripser_parallel(
+            xx.detach().numpy(),
+            maxdim=max(self.homology_dimensions),
+            thresh=self.max_edge_length,
+            coeff=2,
+            metric=self.metric,
+            collapse_edges=self.collapse_edges,
+            n_threads=-1,
+            return_generators=True,
+        )
+
         persistence_pairs = []
 
         for dim in self.homology_dimensions:
             if dim == 0:
-                persistence_pairs += [(0, torch.norm(xx[x[1]]-xx[x[2]]),
-                                      0) for x in output["gens"][dim]]
+                persistence_pairs += [
+                    (0, torch.norm(xx[x[1]] - xx[x[2]]), 0) for x in output["gens"][dim]
+                ]
             else:
-                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]]
+                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, xx: Tensor) -> Tensor:
@@ -318,17 +334,17 @@ def persistence_function(self, xx: Tensor) -> Tensor:
             persistence_array, phi, _ = self._persistence(xx)
             for item in persistence_array:
                 if item[1] != -1:
-                    out += (phi[item[1]]-phi[item[0]])
+                    out += phi[item[1]] - phi[item[0]]
         else:
             persistence_array = self._computing_persistence_with_gph(xx)
             for item in persistence_array:
-                out += item[1]-item[0]
+                out += item[1] - item[0]
         reg = (xx**2).sum()  # regularisation term
-        return -out + self.zeta*reg  # maximise persistence
+        return -out + self.zeta * reg  # maximise persistence
 
-    def sgd(self, xx: Tensor,
-            lr: float = 0.01,
-            n_epochs: int = 5) -> Tuple[Figure, Figure, List[float]]:
+    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``
@@ -368,30 +384,29 @@ def sgd(self, xx: Tensor,
             y = np.concatenate((y, xx.detach().numpy()[:, 1]))
             loss.backward()  # compute gradients and store them in Xp.grad
             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]))
+            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, xx.detach().numpy()[:, 2]))
-                w = np.concatenate((w,
-                                    1/grads.norm(2) *
-                                    grads.numpy()[:, 2]))
+                w = np.concatenate((w, 1 / grads.norm(2) * grads.numpy()[:, 2]))
             except IndexError:
-                z = np.concatenate((z, 0*xx.detach().numpy()[:, 1]))
-                w = np.concatenate((w, 0*grads.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 = Figure(data=go.Cone(
-            x=x,
-            y=y,
-            z=z,
-            u=u,
-            v=v,
-            w=w,
-            sizemode="absolute",
-            sizeref=2,
-            anchor="tip"))
+        fig3d = Figure(
+            data=go.Cone(
+                x=x,
+                y=y,
+                z=z,
+                u=u,
+                v=v,
+                w=w,
+                sizemode="absolute",
+                sizeref=2,
+                anchor="tip",
+            )
+        )
         return fig, fig3d, loss_val
 
 
@@ -429,17 +444,18 @@ def _combinations_with_single(tupl, max_length: int) -> Array:
             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]), )) + \
-            list((max_length-1)*[(tupl[0], tupl[0])])
+        list1 = list(((tupl[0], tupl[0]),)) + list(
+            (max_length - 1) * [(tupl[0], tupl[0])]
+        )
         return np.array(list1)
     temp_list = list(combinations(tupl, 2))
-    list2 = temp_list+((max_length-len(temp_list))*[(tupl[0], tupl[1])])
+    list2 = temp_list + ((max_length - len(temp_list)) * [(tupl[0], tupl[1])])
     return np.array(list2)
 
 
 def comb(nnn: int, kkk: int) -> int:
     """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)
+    kkk = min(kkk, nnn - kkk)
+    numer = reduce(op.mul, range(nnn, nnn - kkk, -1), 1)
+    denom = reduce(op.mul, range(1, kkk + 1), 1)
     return numer // denom
diff --git a/gdeep/utility/optimisation/sam.py b/gdeep/utility/optimisation/sam.py
index bc1f820a..6c364236 100644
--- a/gdeep/utility/optimisation/sam.py
+++ b/gdeep/utility/optimisation/sam.py
@@ -4,8 +4,10 @@
 class MissingClosureError(Exception):
     pass
 
+
 class SAM(torch.optim.Optimizer):
     base_optimizer_cls = None
+
     def __init__(self, params, base_optimizer=None, rho=0.05, adaptive=False, **kwargs):
         assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
 
@@ -23,47 +25,66 @@ def first_step(self, zero_grad=False):
             scale = group["rho"] / (grad_norm + 1e-12)
 
             for p in group["params"]:
-                if p.grad is None: continue
-                e_w = (torch.pow(p, 2) if group["adaptive"] else 1.0) * p.grad * scale.to(p)
+                if p.grad is None:
+                    continue
+                e_w = (
+                    (torch.pow(p, 2) if group["adaptive"] else 1.0)
+                    * p.grad
+                    * scale.to(p)
+                )
                 p.add_(e_w)  # climb to the local maximum "w + e(w)"
                 self.state[p]["e_w"] = e_w
 
-        if zero_grad: self.zero_grad()
+        if zero_grad:
+            self.zero_grad()
 
     @torch.no_grad()
     def second_step(self, zero_grad=False):
         for group in self.param_groups:
             for p in group["params"]:
-                if p.grad is None: continue
+                if p.grad is None:
+                    continue
                 p.sub_(self.state[p]["e_w"])  # get back to "w" from "w + e(w)"
 
         self.base_optimizer.step()  # do the actual "sharpness-aware" update
 
-        if zero_grad: self.zero_grad()
+        if zero_grad:
+            self.zero_grad()
 
     @torch.no_grad()
     def step(self, closure=None):
         if closure is None:
-            raise MissingClosureError("Sharpness Aware Minimization requires closure, but it was not provided")
-        closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass
+            raise MissingClosureError(
+                "Sharpness Aware Minimization requires closure, but it was not provided"
+            )
+        closure = torch.enable_grad()(
+            closure
+        )  # the closure should do a full forward-backward pass
 
         self.first_step(zero_grad=True)
         closure()
         self.second_step()
 
     def _grad_norm(self):
-        shared_device = self.param_groups[0]["params"][0].device  # put everything on the same device, in case of model parallelism
+        shared_device = self.param_groups[0]["params"][
+            0
+        ].device  # put everything on the same device, in case of model parallelism
         norm = torch.norm(
-                    torch.stack([
-                        ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
-                        for group in self.param_groups for p in group["params"]
-                        if p.grad is not None
-                    ]),
-                    p=2
-               )
+            torch.stack(
+                [
+                    ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad)
+                    .norm(p=2)
+                    .to(shared_device)
+                    for group in self.param_groups
+                    for p in group["params"]
+                    if p.grad is not None
+                ]
+            ),
+            p=2,
+        )
         return norm
 
-    
+
 class SAMOptimizer(SAM):
     """This is the class to be used in the pipeline.
 
@@ -78,7 +99,10 @@ class SAMOptimizer(SAM):
         optimizer (torch.optim):
             the optimiser class (not the instance)
     """
-    def __new__(cls, optimizer):
-        return type("SAMOptimizer", (SAM, torch.optim.Optimizer), {"base_optimizer_cls": optimizer})
 
-        
+    def __new__(cls, optimizer):
+        return type(
+            "SAMOptimizer",
+            (SAM, torch.optim.Optimizer),
+            {"base_optimizer_cls": optimizer},
+        )
diff --git a/gdeep/utility/utils.py b/gdeep/utility/utils.py
index 4854bdec..577b500f 100644
--- a/gdeep/utility/utils.py
+++ b/gdeep/utility/utils.py
@@ -15,8 +15,8 @@
 
 
 def _are_compatible(
-        model_dict: Dict[str, Union[nn.Module, str]],
-        dataloaders_dict: Dict[str, Union[List[DataLoader[Tuple[Tensor, Tensor]]], str]]
+    model_dict: Dict[str, Union[nn.Module, str]],
+    dataloaders_dict: Dict[str, Union[List[DataLoader[Tuple[Tensor, Tensor]]], str]],
 ) -> bool:
     """utility function to check the compatibility of a model
     with a set of dataloaders `(dl_tr, dl_val, dl_ts)`
@@ -153,7 +153,7 @@ def _inner_refactor_scalars(list_, cross_validation, k_folds):
 
 def is_notebook() -> bool:
     """Check if the current environment is a notebook
-    
+
     Returns:
         bool:
             True if the environment is a notebook, False otherwise
@@ -172,7 +172,7 @@ def is_notebook() -> bool:
 
 def autoreload_if_notebook() -> None:
     """Autoreload the modules if the environment is a notebook
-    
+
     Returns:
         None
     """
@@ -205,7 +205,7 @@ def get_checksum(file: str, encoding: str = "hex"):
             Path to the file
         encoding:
             string encoding like ``'hex'``
-            
+
     Raises:
         ValueError: if the file does not exist
         ValueError: if the encoding is not supported
@@ -241,7 +241,7 @@ def __enter__(self) -> "KnownWarningSilencer":
     def __exit__(self, type, value, traceback):
         warnings.filterwarnings("default")
 
-        
+
 def flatten_list_of_lists(list_: list) -> list:
     """Flatten a list of lists
     Args:
diff --git a/gdeep/visualisation/__init__.py b/gdeep/visualisation/__init__.py
index 0d1d1d6f..8842ee2b 100644
--- a/gdeep/visualisation/__init__.py
+++ b/gdeep/visualisation/__init__.py
@@ -1,5 +1,7 @@
-from .persistence_activations import persistence_diagrams_of_activations, \
-    _simplified_persistence_of_activations
+from .persistence_activations import (
+    persistence_diagrams_of_activations,
+    _simplified_persistence_of_activations,
+)
 from .plot_compactification import Compactification
 from .utils import plotly2tensor, png2tensor
 from .visualiser import Visualiser
diff --git a/gdeep/visualisation/persistence_activations.py b/gdeep/visualisation/persistence_activations.py
index 5b8fe4b2..4c16608e 100644
--- a/gdeep/visualisation/persistence_activations.py
+++ b/gdeep/visualisation/persistence_activations.py
@@ -30,8 +30,12 @@ def knn_distance_matrix(x: List, k: int = 3):
 
 
 def persistence_diagrams_of_activations(
-        activations_list: List[Tensor], homology_dimensions: Optional[List[int]] = None, k: int = 5,
-        mode: str = "VR", max_edge_length: int = 10) -> List[Any]:
+    activations_list: List[Tensor],
+    homology_dimensions: Optional[List[int]] = None,
+    k: int = 5,
+    mode: str = "VR",
+    max_edge_length: int = 10,
+) -> List[Any]:
     """Returns list of persistence diagrams of the activations of all
     layers of type layer_types
 
@@ -102,20 +106,26 @@ def persistence_diagrams_of_activations(
     return persistence_diagrams
 
 
-def _simplified_persistence_of_activations(inputs: List[Tensor], homology_dimensions: List[int],
-                                           filtration_value: float, **kwargs):
+def _simplified_persistence_of_activations(
+    inputs: List[Tensor],
+    homology_dimensions: List[int],
+    filtration_value: float,
+    **kwargs
+):
     """this method filters out the distances between the input points that are below the
     filtration value and sets to 1 those that are above it. This is useful for a simplified
     computation of homology"""
     vr = VietorisRipsPersistence(
-        homology_dimensions=homology_dimensions,
-        metric="precomputed",
-        **kwargs
+        homology_dimensions=homology_dimensions, metric="precomputed", **kwargs
     )
 
     activations_list_array = _convert_list_of_tensor_to_numpy(inputs)
-    activations_distances = [pairwise_distances(x.reshape((x.shape[0], -1))) for x in activations_list_array]
-    activations_filtered_distances = [(array > filtration_value) * 1. for array in activations_distances]
+    activations_distances = [
+        pairwise_distances(x.reshape((x.shape[0], -1))) for x in activations_list_array
+    ]
+    activations_filtered_distances = [
+        (array > filtration_value) * 1.0 for array in activations_distances
+    ]
 
     simplified_persistence_diagrams = vr.fit_transform(activations_filtered_distances)
     return simplified_persistence_diagrams
diff --git a/gdeep/visualisation/plot_compactification.py b/gdeep/visualisation/plot_compactification.py
index 88fae1bc..99eea249 100644
--- a/gdeep/visualisation/plot_compactification.py
+++ b/gdeep/visualisation/plot_compactification.py
@@ -43,6 +43,7 @@ class Compactification:
             the trained network of which to compute
             the boundary
     """
+
     patches: List
     list_of_pts_in_patches: List
     label_final: List
@@ -75,9 +76,7 @@ def __init__(
         self.check_compute_charts = False
 
     @staticmethod
-    def _transition_to_patch(
-        sample_points_tensor: Tensor, i: int
-    ) -> Tensor:
+    def _transition_to_patch(sample_points_tensor: Tensor, i: int) -> Tensor:
         stacking_list = []
         if i == -1:
             points = sample_points_tensor
diff --git a/gdeep/visualisation/utils.py b/gdeep/visualisation/utils.py
index 59462f2f..58415ef6 100644
--- a/gdeep/visualisation/utils.py
+++ b/gdeep/visualisation/utils.py
@@ -21,12 +21,12 @@ def plotly2tensor(fig) -> torch.Tensor:
     """
     now = str(time.time()).replace(".", "-")
     try:
-        write_image(fig, "deleteme"+now+".jpeg", format="jpeg", engine="orca")
+        write_image(fig, "deleteme" + now + ".jpeg", format="jpeg", engine="orca")
     except ValueError:
-        write_image(fig, "deleteme"+now+".jpeg", format="jpeg")
-    with Image.open("deleteme"+now+".jpeg") as img:
+        write_image(fig, "deleteme" + now + ".jpeg", format="jpeg")
+    with Image.open("deleteme" + now + ".jpeg") as img:
         arr = np.asarray(img).copy()
-    os.remove("deleteme"+now+".jpeg")
+    os.remove("deleteme" + now + ".jpeg")
     return torch.from_numpy(arr)
 
 
diff --git a/gdeep/visualisation/visualiser.py b/gdeep/visualisation/visualiser.py
index e3d3f097..a1a7f20e 100644
--- a/gdeep/visualisation/visualiser.py
+++ b/gdeep/visualisation/visualiser.py
@@ -48,6 +48,7 @@ class Visualiser:
         vs.plot_data_model()
 
     """
+
     persistence_diagrams: Optional[List[Array]] = None
 
     def __init__(self, pipe: Trainer) -> None:
@@ -92,9 +93,13 @@ def plot_3d_dataset(self, n_pts: int = 100) -> None:
             else:
                 features_list = features_list + [img[i] for i in range(len(img))]
             try:
-                labels_list = labels_list + [str(lab[i].item()) for i in range(len(lab))]
+                labels_list = labels_list + [
+                    str(lab[i].item()) for i in range(len(lab))
+                ]
             except ValueError:
-                labels_list = labels_list + ["Label not available" for _ in range(len(img))]
+                labels_list = labels_list + [
+                    "Label not available" for _ in range(len(img))
+                ]
             if len(features_list) >= n_pts:
                 break
         max_number = len(features_list)  # effective number of points
@@ -129,7 +134,9 @@ def plot_3d_dataset(self, n_pts: int = 100) -> None:
 
         self.pipe.writer.flush()  # type: ignore
 
-    def plot_activations(self, batch: Optional[List[Union[List[Tensor], Tensor]]] = None) -> None:
+    def plot_activations(
+        self, batch: Optional[List[Union[List[Tensor], Tensor]]] = None
+    ) -> None:
         """Plot PCA of the activations of all layers of the
         `self.model`.
 
@@ -148,7 +155,7 @@ def plot_activations(self, batch: Optional[List[Union[List[Tensor], Tensor]]] =
         for i, act in enumerate(acts):
             print("Step " + str(i + 1) + "/" + str(len(acts)), end="\r")
             length = act.shape[0]
-            self.pipe.writer.add_embedding(   # type: ignore
+            self.pipe.writer.add_embedding(  # type: ignore
                 act.view(length, -1),
                 metadata=labels,
                 tag="activations_" + str(i),
@@ -156,9 +163,12 @@ def plot_activations(self, batch: Optional[List[Union[List[Tensor], Tensor]]] =
             )
             self.pipe.writer.flush()  # type: ignore
 
-    def plot_persistence_diagrams(self, batch: Optional[List[Tensor]] = None,
-                                  homology_dimensions: Optional[List[int]] = None,
-                                  **kwargs) -> None:  # type: ignore
+    def plot_persistence_diagrams(
+        self,
+        batch: Optional[List[Tensor]] = None,
+        homology_dimensions: Optional[List[int]] = None,
+        **kwargs
+    ) -> None:  # type: ignore
         """Plot a persistence diagrams of the activations
 
         Args:
@@ -180,9 +190,9 @@ def plot_persistence_diagrams(self, batch: Optional[List[Tensor]] = None,
                 inputs, _ = next(iter(self.pipe.dataloaders[0]))
 
             activ = me.get_activations(inputs)
-            self.persistence_diagrams = persistence_diagrams_of_activations(activ,
-                                                                            homology_dimensions=homology_dimensions,
-                                                                            **kwargs)
+            self.persistence_diagrams = persistence_diagrams_of_activations(
+                activ, homology_dimensions=homology_dimensions, **kwargs
+            )
         list_of_dgms = []
         for i, persistence_diagram in enumerate(self.persistence_diagrams):
             plot_persistence_diagram = plot_diagram(persistence_diagram)
@@ -195,7 +205,9 @@ def plot_persistence_diagrams(self, batch: Optional[List[Tensor]] = None,
         )
         self.pipe.writer.flush()  # type: ignore
 
-    def plot_decision_boundary(self, compact: bool = False, n_epochs: int = 100, precision: float = 0.1):
+    def plot_decision_boundary(
+        self, compact: bool = False, n_epochs: int = 100, precision: float = 0.1
+    ):
         """Plot the decision boundary as the intrinsic
         hypersurface defined by ``self.loss_fn == 0.5``
         in the case of 2 classes. The method also works
@@ -249,9 +261,11 @@ def plot_decision_boundary(self, compact: bool = False, n_epochs: int = 100, pre
             return db.cpu(), None, None
 
     def plot_betti_surface_layers(
-            self, homology_dimensions: Optional[List[int]] = None,
-            batch: Optional[Tensor] = None,
-            **kwargs) -> None:  # type: ignore
+        self,
+        homology_dimensions: Optional[List[int]] = None,
+        batch: Optional[Tensor] = None,
+        **kwargs
+    ) -> None:  # type: ignore
         """
         Args:
             homology_dimensions :
@@ -284,17 +298,18 @@ def plot_betti_surface_layers(
         dgms = bc.fit_transform(self.persistence_diagrams)
 
         plots = plot_betti_surfaces(
-            dgms, samplings=bc.samplings_,
-            homology_dimensions=homology_dimensions
+            dgms, samplings=bc.samplings_, homology_dimensions=homology_dimensions
         )
 
         for i in range(len(plots)):
             plots[i].show()
 
     def plot_betti_curves_layers(
-            self, homology_dimensions: Optional[List[int]] = None,
-            batch: Optional[Tensor] = None,
-            **kwargs) -> None:  # type: ignore
+        self,
+        homology_dimensions: Optional[List[int]] = None,
+        batch: Optional[Tensor] = None,
+        **kwargs
+    ) -> None:  # type: ignore
         """
         Args:
             homology_dimensions :
@@ -328,8 +343,9 @@ def plot_betti_curves_layers(
         list_of_plts = []
         for bc_curve in bc_curves:
             plots = plot_betti_curves(
-                bc_curve, samplings=bc.samplings_,
-                homology_dimensions=homology_dimensions
+                bc_curve,
+                samplings=bc.samplings_,
+                homology_dimensions=homology_dimensions,
             )
             img_t = plotly2tensor(plots)
             list_of_plts.append(img_t)
@@ -347,7 +363,7 @@ def plot_interpreter_text(self, interpreter: Interpreter):
             interpreter :
                 this is a ``gdeep.analysis.interpretability``
                 initilised ``Interpreter`` class
-                
+
         Returns:
             matplotlib.figure
         """
@@ -368,14 +384,17 @@ def plot_interpreter_image(self, interpreter: Interpreter):
             interpreter:
                 this is a ``gdeep.analysis.interpretability``
                 initilised ``Interpreter`` class
-                
+
         Returns:
             matplotlib.figure
         """
 
         try:
             attrib = (
-                torch.permute(interpreter.attribution.squeeze().detach(), (1, 2, 0)).detach().cpu().numpy()
+                torch.permute(interpreter.attribution.squeeze().detach(), (1, 2, 0))
+                .detach()
+                .cpu()
+                .numpy()
             )
         except ValueError:
             attrib = torch.permute(
@@ -386,10 +405,17 @@ def plot_interpreter_image(self, interpreter: Interpreter):
                 (1, 2, 0),
             )
             attrib = (
-                torch.stack([attrib, attrib, attrib], dim=2).squeeze(-1).detach().cpu().numpy()
+                torch.stack([attrib, attrib, attrib], dim=2)
+                .squeeze(-1)
+                .detach()
+                .cpu()
+                .numpy()
             )
         img = (
-            torch.permute(interpreter.x.squeeze().detach(), (1, 2, 0)).detach().cpu().numpy()
+            torch.permute(interpreter.x.squeeze().detach(), (1, 2, 0))
+            .detach()
+            .cpu()
+            .numpy()
         )
 
         fig, _ = visualization.visualize_image_attr_multiple(
@@ -410,7 +436,7 @@ def plot_feature_importance(self, interpreter: Interpreter):
             interpreter :
                 this is a ``gdeep.analysis.interpretability``
                 initialised ``Interpreter`` class
-                
+
         Returns:
             matplotlib.figure
         """
@@ -437,9 +463,11 @@ def plot_feature_importance(self, interpreter: Interpreter):
         plt.rc("legend", fontsize=FONT_SIZE - 4)  # fontsize of the legend
         for i in range(len(interpreter.attribution_list)):  # type: ignore
             attribution_sum = interpreter.attribution_list[i].detach().cpu().numpy().sum(0)  # type: ignore
-            attribution_norm_sum = attribution_sum / np.linalg.norm(attribution_sum, ord=1)
+            attribution_norm_sum = attribution_sum / np.linalg.norm(
+                attribution_sum, ord=1
+            )
             r = lambda: random.randint(0, 255)  # noqa
-            color = ('#%02X%02X%02X' % (r(), r(), r()))
+            color = "#%02X%02X%02X" % (r(), r(), r())
             ax.bar(
                 x_axis_data,
                 attribution_norm_sum,
@@ -459,7 +487,9 @@ def plot_feature_importance(self, interpreter: Interpreter):
         self.pipe.writer.add_figure("Feature Importance", fig)  # type: ignore
         return plt
 
-    def plot_attribution(self, interpreter: Interpreter, **kwargs) -> Tuple[plt.Figure, plt.Figure]:
+    def plot_attribution(
+        self, interpreter: Interpreter, **kwargs
+    ) -> Tuple[plt.Figure, plt.Figure]:
         """this method generically plots the attribution of
         the interpreter
 
@@ -482,8 +512,13 @@ def plot_attribution(self, interpreter: Interpreter, **kwargs) -> Tuple[plt.Figu
         self.pipe.writer.add_figure("Generic attribution of x", fig2)  # type: ignore
         return fig1, fig2
 
-    def plot_self_attention(self, attention_tensor: List[Tensor], tokens_x: Optional[List[str]] = None,
-                            tokens_y: Optional[List[str]] = None, **kwargs) -> plt.Figure:
+    def plot_self_attention(
+        self,
+        attention_tensor: List[Tensor],
+        tokens_x: Optional[List[str]] = None,
+        tokens_y: Optional[List[str]] = None,
+        **kwargs
+    ) -> plt.Figure:
         """This functions plots the self-attention layer of a transformer.
 
         Args:
@@ -500,16 +535,16 @@ def plot_self_attention(self, attention_tensor: List[Tensor], tokens_x: Optional
         Returns:
              Figure:
                  matplotlib pyplot
-            """
+        """
         fig = plt.figure(**kwargs)
 
         for idx, scores in enumerate(attention_tensor):
             scores_np = Visualiser._adjust_tensors_to_plot(scores)
             ax = fig.add_subplot(4, 3, idx + 1)
             # append the attention weights
-            im = ax.imshow(scores_np, cmap='viridis')
+            im = ax.imshow(scores_np, cmap="viridis")
 
-            fontdict = {'fontsize': 10}
+            fontdict = {"fontsize": 10}
             if tokens_x:
                 ax.set_xticks(range(len(tokens_x)))
                 ax.set_xticklabels(tokens_x, fontdict=fontdict, rotation=90)
@@ -517,7 +552,7 @@ def plot_self_attention(self, attention_tensor: List[Tensor], tokens_x: Optional
                 ax.set_yticks(range(len(tokens_y)))
                 ax.set_yticklabels(tokens_y, fontdict=fontdict)
 
-            ax.set_xlabel('Head {}'.format(idx + 1))
+            ax.set_xlabel("Head {}".format(idx + 1))
 
             fig.colorbar(im, fraction=0.046, pad=0.04)
         plt.show()
@@ -525,7 +560,9 @@ def plot_self_attention(self, attention_tensor: List[Tensor], tokens_x: Optional
         self.pipe.writer.flush()  # type: ignore
         return fig
 
-    def plot_attributions_persistence_diagrams(self, interpreter: Interpreter, **kwargs) -> plt.Figure:
+    def plot_attributions_persistence_diagrams(
+        self, interpreter: Interpreter, **kwargs
+    ) -> plt.Figure:
         """this method allows the plot, on top of the persistence
         diagram, of the attribution values. For example, this would
         be the method to call when you run saliency maps on the
@@ -551,7 +588,7 @@ def plot_attributions_persistence_diagrams(self, interpreter: Interpreter, **kwa
         fig = plt.figure(figsize=(10, 10))
         plt.scatter(x=pd_x.T[0], y=pd_x.T[1], c=yy, **kwargs)
         line = np.arange(0, max(pd_x.T[1]) * 1.1, 0.01)
-        plt.plot(line, line, '--', color="black")
+        plt.plot(line, line, "--", color="black")
         plt.xlabel("Birth")
         plt.ylabel("Death")
 
@@ -561,10 +598,12 @@ def plot_attributions_persistence_diagrams(self, interpreter: Interpreter, **kwa
         return fig
 
     def plot_betti_numbers_layers(
-            self, homology_dimensions: Optional[List[int]] = None,
-            filtration_value: float = 1.,
-            batch: Optional[Tensor] = None,
-            **kwargs) -> plt.Figure:  # type: ignore
+        self,
+        homology_dimensions: Optional[List[int]] = None,
+        filtration_value: float = 1.0,
+        batch: Optional[Tensor] = None,
+        **kwargs
+    ) -> plt.Figure:  # type: ignore
         """
         Args:
             homology_dimensions :
@@ -592,16 +631,19 @@ def plot_betti_numbers_layers(
         else:
             inputs = next(iter(self.pipe.dataloaders[0]))  # type: ignore
 
-        simplified_persistence_diagrams = _simplified_persistence_of_activations(me.get_activations(inputs),
-                                                                                 homology_dimensions,
-                                                                                 filtration_value, ** kwargs)
+        simplified_persistence_diagrams = _simplified_persistence_of_activations(
+            me.get_activations(inputs), homology_dimensions, filtration_value, **kwargs
+        )
 
         betti_numbers = []
 
         for dgm in simplified_persistence_diagrams:
             xx = dgm[:, 2]
             betti_numbers.append(
-                [len(xx[(dgm[:, 2] == dim) * (dgm[:, 1] == 1.)]) for dim in homology_dimensions]
+                [
+                    len(xx[(dgm[:, 2] == dim) * (dgm[:, 1] == 1.0)])
+                    for dim in homology_dimensions
+                ]
             )
         betti_array = np.array(betti_numbers).T
 
diff --git a/kubernetes/dockerfiles/jupyter-lab/Dockerfile b/kubernetes/dockerfiles/jupyter-lab/Dockerfile
index 810c570b..9ad908d7 100644
--- a/kubernetes/dockerfiles/jupyter-lab/Dockerfile
+++ b/kubernetes/dockerfiles/jupyter-lab/Dockerfile
@@ -3,7 +3,7 @@ FROM ubuntu:20.04
 RUN apt-get update && apt-get install -y sudo openssh-server vim
 RUN apt-get install -y supervisor
 RUN mkdir /var/run/sshd
-RUN useradd -rm -d /home/ubuntu -s /bin/bash -g root -G sudo -u 1000 matteo 
+RUN useradd -rm -d /home/ubuntu -s /bin/bash -g root -G sudo -u 1000 matteo
 RUN echo 'matteo:i_like_giotto' | chpasswd
 RUN sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config
 
@@ -30,7 +30,7 @@ RUN pip install mysqlclient
 RUN apt-get install -y git
 
 # for multiprocesses
-ADD kubernetes/dockerfiles/jupyter-lab/supervisord.conf supervisord.conf 
+ADD kubernetes/dockerfiles/jupyter-lab/supervisord.conf supervisord.conf
 
 # expose the port for jupyter and ssh
 EXPOSE 8888 22
@@ -39,4 +39,4 @@ EXPOSE 8888 22
 RUN mkdir giotto-deep
 VOLUME ["/giotto-deep"]
 
-CMD ["supervisord", "-c", "supervisord.conf"]
\ No newline at end of file
+CMD ["supervisord", "-c", "supervisord.conf"]
diff --git a/kubernetes/dockerfiles/jupyter-lab/supervisord.conf b/kubernetes/dockerfiles/jupyter-lab/supervisord.conf
index 4113c0bd..a079915b 100644
--- a/kubernetes/dockerfiles/jupyter-lab/supervisord.conf
+++ b/kubernetes/dockerfiles/jupyter-lab/supervisord.conf
@@ -13,4 +13,4 @@ directory=/
 command=jupyter-lab --allow-root --ip=0.0.0.0 --ServerApp.token='abcd'
 autostart=true
 autorestart=true
-redirect_stderr=true
\ No newline at end of file
+redirect_stderr=true
diff --git a/kubernetes/dockerfiles/rq-workers/Dockerfile b/kubernetes/dockerfiles/rq-workers/Dockerfile
index 8d29ab10..018d33fa 100644
--- a/kubernetes/dockerfiles/rq-workers/Dockerfile
+++ b/kubernetes/dockerfiles/rq-workers/Dockerfile
@@ -2,4 +2,4 @@ FROM matteocao/giotto-deep:gdeep-lab
 
 RUN ["pip3", "install", "giotto-deep"]
 
-CMD ["rq", "worker", "--url", "redis://redis-service", "high", "default", "low"]
\ No newline at end of file
+CMD ["rq", "worker", "--url", "redis://redis-service", "high", "default", "low"]
diff --git a/kubernetes/dockerfiles/rq-workers/supervisord.conf b/kubernetes/dockerfiles/rq-workers/supervisord.conf
index 4113c0bd..a079915b 100644
--- a/kubernetes/dockerfiles/rq-workers/supervisord.conf
+++ b/kubernetes/dockerfiles/rq-workers/supervisord.conf
@@ -13,4 +13,4 @@ directory=/
 command=jupyter-lab --allow-root --ip=0.0.0.0 --ServerApp.token='abcd'
 autostart=true
 autorestart=true
-redirect_stderr=true
\ No newline at end of file
+redirect_stderr=true
diff --git a/kubernetes/examples/parallel_hpo.py b/kubernetes/examples/parallel_hpo.py
index fd896794..959aec3f 100644
--- a/kubernetes/examples/parallel_hpo.py
+++ b/kubernetes/examples/parallel_hpo.py
@@ -1,4 +1,3 @@
-
 import socket
 from datetime import datetime
 import os
@@ -28,8 +27,10 @@ def hpo_parallel(usr: str, psw: str, host: str, study_name: str) -> None:
             the name of the optuna study
     """
 
-    current_time: str = datetime.now().strftime('%b%d_%H-%M-%S')
-    log_dir = os.path.join('giotto-deep', 'runs', current_time + '_' + socket.gethostname())
+    current_time: str = datetime.now().strftime("%b%d_%H-%M-%S")
+    log_dir = os.path.join(
+        "giotto-deep", "runs", current_time + "_" + socket.gethostname()
+    )
     writer = GiottoSummaryWriter(log_dir=log_dir)
 
     bd = DatasetBuilder(name="DoubleTori")
@@ -64,8 +65,13 @@ def forward(self, x):
     # initialise the SAM optimiser
     optim = SGD  # this is a class, not an instance!
 
-    search = HyperParameterOptimization(pipe, "loss", 20, study_name=study_name,
-                                        db_url="mysql+mysqldb://" + usr + ":" + psw + "@" + host + ":3306/example")
+    search = HyperParameterOptimization(
+        pipe,
+        "loss",
+        20,
+        study_name=study_name,
+        db_url="mysql+mysqldb://" + usr + ":" + psw + "@" + host + ":3306/example",
+    )
     # dictionaries of hyperparameters
     optimizers_params = {"lr": [0.001, 0.01]}
     dataloaders_params = {"batch_size": [32, 64, 16]}
diff --git a/kubernetes/examples/utils.py b/kubernetes/examples/utils.py
index 32d4ffca..595fce8f 100644
--- a/kubernetes/examples/utils.py
+++ b/kubernetes/examples/utils.py
@@ -17,11 +17,8 @@ def connect_to_mysql(usr: str, psw: str, host: str) -> None:
             IP of the host. Port is 3306
     """
     import mysql.connector
-    mydb = mysql.connector.connect(
-        host=host,
-        user=usr,
-        password=psw
-    )
+
+    mydb = mysql.connector.connect(host=host, user=usr, password=psw)
 
     print(mydb)
     mycursor = mydb.cursor()
@@ -64,4 +61,3 @@ def run_hpo_in_parallel(q: Queue, fnc: Callable, args: List[Any], number: int) -
     """
     for _ in range(number):
         q.enqueue(fnc, args=args, retry=Retry(max=3))
-
diff --git a/kubernetes/gdeep-lab.yaml b/kubernetes/gdeep-lab.yaml
index 9c87ca09..873799b2 100644
--- a/kubernetes/gdeep-lab.yaml
+++ b/kubernetes/gdeep-lab.yaml
@@ -50,4 +50,4 @@ spec:
     - protocol: TCP
       port: 22
       targetPort: 22  #must match container port
-      name: ssh
\ No newline at end of file
+      name: ssh
diff --git a/kubernetes/lab-ingress.yaml b/kubernetes/lab-ingress.yaml
index e5a08105..1a88a601 100644
--- a/kubernetes/lab-ingress.yaml
+++ b/kubernetes/lab-ingress.yaml
@@ -22,7 +22,7 @@ spec:
             backend:
               service:
                 name: gdeep-lab-service
-                port: 
+                port:
                   number: 22
           - path: /plot
             pathType: Prefix
@@ -30,4 +30,4 @@ spec:
               service:
                 name: gdeep-tensorboard-service
                 port:
-                  number: 8080
\ No newline at end of file
+                  number: 8080
diff --git a/kubernetes/mysql.yaml b/kubernetes/mysql.yaml
index 6c901937..95bac1bf 100644
--- a/kubernetes/mysql.yaml
+++ b/kubernetes/mysql.yaml
@@ -22,7 +22,7 @@ spec:
         - containerPort: 3306
         env:
         - name: MYSQL_ROOT_PASSWORD
-          valueFrom: 
+          valueFrom:
             secretKeyRef:
               name: mysql-secret
               key: mysql-root-password
diff --git a/kubernetes/namespace-main.json b/kubernetes/namespace-main.json
index 8a35ec32..ed2c3acb 100644
--- a/kubernetes/namespace-main.json
+++ b/kubernetes/namespace-main.json
@@ -7,4 +7,4 @@
       "name": "main"
     }
   }
-}
\ No newline at end of file
+}
diff --git a/kubernetes/pv-volume-gcp.yaml b/kubernetes/pv-volume-gcp.yaml
index 6851bb85..ac0ad8e0 100644
--- a/kubernetes/pv-volume-gcp.yaml
+++ b/kubernetes/pv-volume-gcp.yaml
@@ -19,4 +19,4 @@ spec:
           operator: In
           values:
           - us-central1-a
-          - us-central1-b
\ No newline at end of file
+          - us-central1-b
diff --git a/kubernetes/pv-volume.yaml b/kubernetes/pv-volume.yaml
index 3d473743..4b82c572 100644
--- a/kubernetes/pv-volume.yaml
+++ b/kubernetes/pv-volume.yaml
@@ -11,4 +11,4 @@ spec:
   accessModes:
     - ReadWriteOnce
   hostPath:
-    path: "/data"
\ No newline at end of file
+    path: "/data"
diff --git a/kubernetes/readme.md b/kubernetes/readme.md
index 679cc69e..5700d213 100644
--- a/kubernetes/readme.md
+++ b/kubernetes/readme.md
@@ -20,14 +20,14 @@ docker push matteocao/giotto-deep:gdeep-lab
 All the config files should work out of the blue, just follow these steps:
  1. Move all the files to the kubernetes master server
  2. Set up a persistent machine for volume (no need if using `minikube` on local for testing):
- 
+
 ```
 gcloud compute disks create --size=200GB my-data-disk --region us-central1 --replica-zones us-central1-a,us-central1-b
 ```
 
  3. Create namespace `kubectl create namespace main`
  4. Run the following commands, in this order:
- 
+
 ```
 kubectl create -f namespace-main.json
 kubectl apply -f pv-volume.yaml  # or kubectl apply -f pv-volume-gcp.yaml if on GCP
@@ -38,12 +38,12 @@ kubectl apply -f tensorboard.yaml
 kubectl apply -f gdeep-lab.yaml
 kubectl apply -f redis.yaml
 kubectl apply -f rq-worker.yaml  # to run the workers; if you change the location of the parallel_hpo.py file, make sure to change the --path here!
-kubectl apply -f lab-ingress.yaml 
+kubectl apply -f lab-ingress.yaml
 ```
- 
+
  5. Connect to  `<ingress_external_ip>/one:8888` to access jupyter lab. The token is `abcd`. Alternatively, on minikube, please run `minikube service gdeep-lab-service --url -n main` and you will see the URL at the terminal
  6. Connect to `<ingress_external_ip>/ssh_one` if you want to access via ssh. For ssh:
- 
+
 ```
 user: matteo
 psw: i_like_giotto
@@ -56,7 +56,7 @@ kubectl scale deployment gdeep-rq-worker-deployment --replicas=0
 ```
 
  8. Access the jupyter lab service and in the giotto-deep folder (the volume!) you can `git clone https://github.com/giotto-ai/giotto-deep.git` and also `pip install giotto-deep`.
- 9. Run the `giotto-deep/kubernetes/examples/parallel_hpo.ipynb` notebook and see the parallelisation happening! 
+ 9. Run the `giotto-deep/kubernetes/examples/parallel_hpo.ipynb` notebook and see the parallelisation happening!
  10. To check the results, use the tensorboard service: on minikube you can do `minikube service gdeep-tensorboard-service --url -n main`, while on K8 you can simply go to `<ingress_external_ip>/one:8080`
 
 # Set-up Minikube
@@ -82,7 +82,7 @@ or
 kubectl delete service gdeep-lab-service-1 --namespace main
 ```
 
-Afterwards, the cluster is up and running, except for the ingress, which `minikube` does not support. 
+Afterwards, the cluster is up and running, except for the ingress, which `minikube` does not support.
 To check the cluster status simply run:
 
 ```
@@ -136,9 +136,9 @@ minikube delete --all
 ```
 
 ## Working on different machines
- 
+
 You can scale-up workers by changing the number of replicas of the RQ workers deployment (see point 7 above). You can add, to each worker, GPUs or other computing hardware.
- 
+
 # Distributing HPO
 
 To run hyperparameters optimisations in parallel, we need the support of a `mysql` database. To connect to such database, you need to run the python script follow the instruction on `parallel_hpo.ipynb` from the jupyter-lab instance. The script connects optuna to the database. You will need to set up:
@@ -148,9 +148,9 @@ To run hyperparameters optimisations in parallel, we need the support of a `mysq
 You can find the parameters `<mysql_internal_ip>` and `<redis_internal_ip>` by looking at your cluster with `kubectl get all --namespace main` and checking the `ClusterIP` of the `service/mysql-service` and `service/redis-service`.
 Finally, to parallelise the optuna computations, recall to set the parameters in the `HyperParameterOptimisation.__init__`:
  ```
-study_name="distributed-example", 
+study_name="distributed-example",
 db_url="mysql+mysqldb://root:password@<mysql_internal_ip>:3306/example",
  ```
- 
+
 Note that such parameters are already set up in the functions provided: you just need to initialise the variables `USER` and `PSW`
 where you have only to change `<mysql_internal_ip>` with the ClusterIP found before.
diff --git a/kubernetes/rq-worker.yaml b/kubernetes/rq-worker.yaml
index d9d7e245..48ceee88 100644
--- a/kubernetes/rq-worker.yaml
+++ b/kubernetes/rq-worker.yaml
@@ -25,7 +25,7 @@ spec:
         volumeMounts:
         - mountPath: "/giotto-deep"
           name: task-pv-storage
-        command: ["rq"] 
+        command: ["rq"]
         args: ["worker", "--url", "redis://redis-service", "--path", "/giotto-deep/giotto-deep/kubernetes/examples", "high", "default", "low"]
 ---
 apiVersion: v1
diff --git a/requirements.txt b/requirements.txt
index 1e420292..644fe836 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -42,4 +42,5 @@ typing_extensions; python_version == '3.7'
 torch-geometric
 torch-sparse
 torch-scatter
-gudhi
\ No newline at end of file
+gudhi
+pre-commit
diff --git a/setup.py b/setup.py
index 94d10690..10b546e7 100755
--- a/setup.py
+++ b/setup.py
@@ -8,7 +8,7 @@
 
 
 version_file = os.path.join("gdeep", "_version.py")
-__version__ = '0.0.0'  # initialise the variable
+__version__ = "0.0.0"  # initialise the variable
 with open(version_file) as f:
     exec(f.read())
 
@@ -25,52 +25,57 @@
 URL = "https://github.com/giotto-ai/giotto-deep"
 LICENSE = "GNU AGPLv3"
 VERSION = __version__  # type: ignore
-DOWNLOAD_URL = "https://github.com/giotto-ai/giotto-deep/tarball/v"+VERSION
-CLASSIFIERS = ["Intended Audience :: Science/Research",
-               "Intended Audience :: Developers",
-               "License :: OSI Approved",
-               "Programming Language :: Python",
-               "Topic :: Software Development",
-               "Topic :: Scientific/Engineering",
-               "Operating System :: Microsoft :: Windows",
-               "Operating System :: POSIX",
-               "Operating System :: Unix",
-               "Operating System :: MacOS",
-               "Programming Language :: Python :: 3.7",
-               "Programming Language :: Python :: 3.8",
-               "Programming Language :: Python :: 3.9"]
-KEYWORDS = "deep learning, topological data analysis, persistent " \
-           "homology, persistence diagrams"
+DOWNLOAD_URL = "https://github.com/giotto-ai/giotto-deep/tarball/v" + VERSION
+CLASSIFIERS = [
+    "Intended Audience :: Science/Research",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved",
+    "Programming Language :: Python",
+    "Topic :: Software Development",
+    "Topic :: Scientific/Engineering",
+    "Operating System :: Microsoft :: Windows",
+    "Operating System :: POSIX",
+    "Operating System :: Unix",
+    "Operating System :: MacOS",
+    "Programming Language :: Python :: 3.7",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+]
+KEYWORDS = (
+    "deep learning, topological data analysis, persistent "
+    "homology, persistence diagrams"
+)
 INSTALL_REQUIRES = requirements
-EXTRAS_REQUIRE = {"tests": ["pytest",
-                            "flake8",
-                            "mypy"],
-                  "doc": ["openml",
-                          "sphinx",
-                          "nbconvert",
-                          "sphinx-issues",
-                          "sphinx_rtd_theme",
-                          "numpydoc"],
-                  "examples": ["jupyter",
-                               "pandas",
-                               "openml",
-                               "plotly"]}
+EXTRAS_REQUIRE = {
+    "tests": ["pytest", "flake8", "mypy"],
+    "doc": [
+        "openml",
+        "sphinx",
+        "nbconvert",
+        "sphinx-issues",
+        "sphinx_rtd_theme",
+        "numpydoc",
+    ],
+    "examples": ["jupyter", "pandas", "openml", "plotly"],
+}
 
-setup(name=DISTNAME,
-      maintainer=MAINTAINER,
-      maintainer_email=MAINTAINER_EMAIL,
-      description=DESCRIPTION,
-      license=LICENSE,
-      url=URL,
-      version=VERSION,
-      download_url=DOWNLOAD_URL,
-      long_description=LONG_DESCRIPTION,
-      long_description_content_type=LONG_DESCRIPTION_TYPE,
-      zip_safe=False,
-      classifiers=CLASSIFIERS,
-      packages=find_packages(),
-      keywords=KEYWORDS,
-      dependency_links=["https://download.pytorch.org/whl/torch_stable.html"],
-      find_links=["https://download.pytorch.org/whl/torch_stable.html"],
-      install_requires=INSTALL_REQUIRES,
-      extras_require=EXTRAS_REQUIRE)
+setup(
+    name=DISTNAME,
+    maintainer=MAINTAINER,
+    maintainer_email=MAINTAINER_EMAIL,
+    description=DESCRIPTION,
+    license=LICENSE,
+    url=URL,
+    version=VERSION,
+    download_url=DOWNLOAD_URL,
+    long_description=LONG_DESCRIPTION,
+    long_description_content_type=LONG_DESCRIPTION_TYPE,
+    zip_safe=False,
+    classifiers=CLASSIFIERS,
+    packages=find_packages(),
+    keywords=KEYWORDS,
+    dependency_links=["https://download.pytorch.org/whl/torch_stable.html"],
+    find_links=["https://download.pytorch.org/whl/torch_stable.html"],
+    install_requires=INSTALL_REQUIRES,
+    extras_require=EXTRAS_REQUIRE,
+)


Test patch:
diff --git a/gdeep/analysis/decision_boundary/tests/test_compute_topology.py b/gdeep/analysis/decision_boundary/tests/test_compute_topology.py
index cae94890..32ddcbe0 100644
--- a/gdeep/analysis/decision_boundary/tests/test_compute_topology.py
+++ b/gdeep/analysis/decision_boundary/tests/test_compute_topology.py
@@ -5,8 +5,7 @@
 
 
 def test_sample_points_uniformly():
-    sample_points = UniformlySampledPoint([(2, 3), (1, 2)],
-                                          n_samples=1000)
+    sample_points = UniformlySampledPoint([(2, 3), (1, 2)], n_samples=1000)
 
     print(sample_points().shape)
     assert sample_points().shape == (1000, 2)
diff --git a/gdeep/analysis/decision_boundary/tests/test_decision_boundary_calculator.py b/gdeep/analysis/decision_boundary/tests/test_decision_boundary_calculator.py
index 462b219d..9461bbb6 100644
--- a/gdeep/analysis/decision_boundary/tests/test_decision_boundary_calculator.py
+++ b/gdeep/analysis/decision_boundary/tests/test_decision_boundary_calculator.py
@@ -13,9 +13,9 @@ def test_gfdbc_2_dim():
     circle_detect_nn = FFNet([2, 2]).to(DEVICE)
 
     g = GradientFlowDecisionBoundaryCalculator(
-            model=circle_detect_nn,
-            initial_points=torch.rand((100, 2)),
-            optimizer=lambda params: torch.optim.Adam(params)
+        model=circle_detect_nn,
+        initial_points=torch.rand((100, 2)),
+        optimizer=lambda params: torch.optim.Adam(params),
     )
     for i in range(100):
         g.step()
@@ -26,15 +26,14 @@ def test_gfdbc_multiclass():
 
     circle_nn_3d = FFNet([3, 3]).to(DEVICE)
 
-    g = GradientFlowDecisionBoundaryCalculator(model=
-                                               circle_nn_3d,
-                                               initial_points=
-                                               torch.rand((100, 3)),
-                                               optimizer=
-                                               lambda params:
-                                               torch.optim.Adam(params))
+    g = GradientFlowDecisionBoundaryCalculator(
+        model=circle_nn_3d,
+        initial_points=torch.rand((100, 3)),
+        optimizer=lambda params: torch.optim.Adam(params),
+    )
     for i in range(100):
         g.step()
     assert g.get_decision_boundary().size() == torch.Size([100, 3])
 
+
 # TODO: Check if significant number of points lie close to the decision boundary
diff --git a/gdeep/analysis/interpretability/tests/test_captum.py b/gdeep/analysis/interpretability/tests/test_captum.py
index a0a70e5a..52b90cf8 100644
--- a/gdeep/analysis/interpretability/tests/test_captum.py
+++ b/gdeep/analysis/interpretability/tests/test_captum.py
@@ -16,8 +16,7 @@
 
 # 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)
+bd = DatasetBuilder(name="AG_NEWS", convert_to_map_dataset=True)
 ds_tr_str, ds_val_str, ds_ts_str = bd.build()
 ptd = TokenizerTextClassification()
 
@@ -29,16 +28,20 @@
 train_indices = list(range(64 * 10))
 test_indices = list(range(64 * 5))
 
-dl_tr2, dl_ts2, _ = DataLoaderBuilder((transformed_textds,  # type: ignore
-                                       transformed_textts)).build(({"batch_size": 16,  # type: ignore
-                                                                    "sampler": SubsetRandomSampler(train_indices)},
-                                                                   {"batch_size": 16,
-                                                                    "sampler": SubsetRandomSampler(test_indices)}
-                                                                   ))
+dl_tr2, dl_ts2, _ = DataLoaderBuilder(
+    (transformed_textds, transformed_textts)  # type: ignore
+).build(
+    (
+        {
+            "batch_size": 16,  # type: ignore
+            "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)
@@ -67,16 +70,16 @@ def forward(self, text):
 
 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,  # type: ignore
-                         ptd.tokenizer,  # type: ignore
-                         layer=pipe.model.embedding,  # type: ignore
-                         n_steps=500,
-                         return_convergence_delta=True
-                         )
+    inter = Interpreter(pipe.model, method="LayerIntegratedGradients")
+
+    inter.interpret_text(
+        "I am writing about money and business",
+        0,
+        ptd.vocabulary,  # type: ignore
+        ptd.tokenizer,  # type: ignore
+        layer=pipe.model.embedding,  # type: ignore
+        n_steps=500,
+        return_convergence_delta=True,
+    )
 
     vs.plot_interpreter_text(inter)
diff --git a/gdeep/data/datasets/tests/__init__.py b/gdeep/data/datasets/tests/__init__.py
index 66247349..5c84a195 100644
--- a/gdeep/data/datasets/tests/__init__.py
+++ b/gdeep/data/datasets/tests/__init__.py
@@ -1,1 +1,1 @@
-#from .data_cloud import DataCloud
\ No newline at end of file
+# from .data_cloud import DataCloud
diff --git a/gdeep/data/datasets/tests/test_data_cloud.py b/gdeep/data/datasets/tests/test_data_cloud.py
index 7d505cfe..a8d18250 100644
--- a/gdeep/data/datasets/tests/test_data_cloud.py
+++ b/gdeep/data/datasets/tests/test_data_cloud.py
@@ -23,12 +23,13 @@
     # Check if the credentials are valid and if the bucket can be accessed
     client = storage.Client()
     if Bucket(client, DATASET_BUCKET_NAME).exists():
+
         def test_download():
             """Test download of sample data from bucket"""
             data_cloud = _DataCloud(use_public_access=False)
             file_name = "giotto-deep-big.png"
             data_cloud.download_file(file_name)
-            
+
             # check if correct extension is raised when trying to download
             # non-existing file
             non_existing_file_name: str = "giotto-deep-bigs.png"
@@ -40,112 +41,108 @@ def test_download():
             assert exists(file_path)
 
             # check if downloaded file is correct
-            assert "d4b12b2dc2bc199831ba803431184fcb" == \
-                get_checksum(file_path)
-                
+            assert "d4b12b2dc2bc199831ba803431184fcb" == get_checksum(file_path)
+
             remove(join(data_cloud.download_directory, file_name))
-                
+
         def test_upload():
             """Test upload of sample file to bucket."""
             data_cloud = _DataCloud(use_public_access=False)
-            
+
             # create temporary file to upload to bucket
             sample_file_name = "tmp.txt"
-            sample_text = ("Create a new tmp file!" 
-                            + str(random.randint(0, 1_000)))
-            
+            sample_text = "Create a new tmp file!" + str(random.randint(0, 1_000))
+
             if exists(sample_file_name):
                 remove(sample_file_name)
-            with open(sample_file_name, 'w') as f:
+            with open(sample_file_name, "w") as f:
                 f.write(sample_text)
-            
+
             assert exists(sample_file_name)
-            
+
             # upload sample file to bucket
             data_cloud.upload_file(sample_file_name)
-            
+
             # check if uploading to an already existing file raises exception
             with pytest.raises(RuntimeError):
                 data_cloud.upload_file(sample_file_name)
-                
+
             # delete local temporary file
             remove(sample_file_name)
-            
+
             data_cloud.download_file(sample_file_name)
-            
+
             data_cloud.delete_blob(sample_file_name)
-            
+
             # check if downloaded file exists
             file_path = join(data_cloud.download_directory, sample_file_name)
             assert exists(file_path)
-            
-            with open(file_path, 'r') as f:
+
+            with open(file_path, "r") as f:
                 assert f.read() == sample_text
-            
+
             remove(file_path)
 
         def test_upload_folder():
-            """Test the upload of a folder to bucket and download the 
+            """Test the upload of a folder to bucket and download the
             folder."""
             data_cloud = _DataCloud(use_public_access=False)
-            
+
             # create temporary folder structure and temporary file to upload
             # to bucket
             # tmp: tmp.txt
             # |- sub_tmp_1: tmp1.txt
             # |- sub_tmp_2: tmp2.txt
             #     |- sub_tmp_2_2: tmp2_2.txt
-            if exists('tmp'):
-                rmtree('tmp')
-            
+            if exists("tmp"):
+                rmtree("tmp")
+
             tmp_files = []
-            
+
             sample_dir = "tmp"
             makedirs(sample_dir)
             tmp_files.append(join(sample_dir, "tmp.txt"))
-            
-            sub_1_sample_dir = join(sample_dir, 'sub_tmp_1')
+
+            sub_1_sample_dir = join(sample_dir, "sub_tmp_1")
             makedirs(sub_1_sample_dir)
             tmp_files.append(join(sub_1_sample_dir, "tmp1.txt"))
-            
-            sub_2_sample_dir = join(sample_dir, 'sub_tmp_2')
+
+            sub_2_sample_dir = join(sample_dir, "sub_tmp_2")
             makedirs(sub_2_sample_dir)
             tmp_files.append(join(sub_2_sample_dir, "tmp2.txt"))
-            
-            sub_2_2_sample_dir = join(sub_2_sample_dir, 'sub_tmp_2_2')
+
+            sub_2_2_sample_dir = join(sub_2_sample_dir, "sub_tmp_2_2")
             makedirs(sub_2_2_sample_dir)
             tmp_files.append(join(sub_2_2_sample_dir, "tmp2_2.txt"))
 
-            
             sample_texts = {}
-            
+
             for file in tmp_files:
                 if exists(file):
                     remove(file)
-                sample_text = ("Create a new tmp file! " + 
-                                str(random.randint(0, 1_000)))
+                sample_text = "Create a new tmp file! " + str(random.randint(0, 1_000))
                 sample_texts[file] = sample_text
-                with open(file, 'w') as f:
+                with open(file, "w") as f:
                     f.write(sample_text)
                 assert exists(file)
-            
+
             # upload sample file to bucket
             data_cloud.upload_folder(sample_dir, "tmp")
-            
+
             # delete local tmp folder
             rmtree(sample_dir)
-            
+
             # download folder to local
             data_cloud.download_folder(sample_dir)
-            
+
             # delete folder in bucket
             data_cloud.delete_blobs("tmp")
-            
+
             # check if downloaded folder is correct
-            
+
             for file in sample_texts.keys():
-                with open(join(data_cloud.download_directory, file), 'r') as f:
+                with open(join(data_cloud.download_directory, file), "r") as f:
                     assert f.read() == sample_texts[file]
-            
+
             # delete local tmp folder
             rmtree(join(data_cloud.download_directory, sample_dir))
diff --git a/gdeep/data/datasets/tests/test_dataset_cloud.py b/gdeep/data/datasets/tests/test_dataset_cloud.py
index 413d332a..930c1018 100644
--- a/gdeep/data/datasets/tests/test_dataset_cloud.py
+++ b/gdeep/data/datasets/tests/test_dataset_cloud.py
@@ -20,90 +20,86 @@ def test_public_access():
     # Download a small dataset from Google Cloud Storage
     dataset = "SmallDataset"
     download_directory = join("examples", "data", "DatasetCloud", "Tmp")
-    
+
     # Remove download directory recursively if it exists
     if exists(download_directory):
         rmtree(download_directory)
-    
+
     # Create download directory
     os.makedirs(download_directory, exist_ok=False)
-    
-    
-    dataset_cloud = DatasetCloud(dataset,
-                                 download_directory=download_directory,
-                                 use_public_access=True)
+
+    dataset_cloud = DatasetCloud(
+        dataset, download_directory=download_directory, use_public_access=True
+    )
     dataset_cloud.download()
-    
+
     # Check if the downloaded files (metadata.json, data.json, labels.json)
     # are correct
-    checksums = {"data.pt": "2ef68a718e29134cbcbf46c9592f6168",
-                 "labels.pt": "d71992425033c6bf449d175db146a423",
-                 }
-    
+    checksums = {
+        "data.pt": "2ef68a718e29134cbcbf46c9592f6168",
+        "labels.pt": "d71992425033c6bf449d175db146a423",
+    }
+
     for file in checksums.keys():
-        assert get_checksum(join(download_directory,
-                                       dataset,
-                                       file)) == \
-            checksums[file], "File {} is corrupted.".format(file)
+        assert (
+            get_checksum(join(download_directory, dataset, file)) == checksums[file]
+        ), "File {} is corrupted.".format(file)
     # Recursively remove download directory
     rmtree(download_directory)
 
 
-    
-           
 def test_get_dataset_list():
     # Download directory will not be used as well ass the dataset
     # It's only used for initialization of the DatasetCloud object
     download_directory = ""
-    dataset_cloud = DatasetCloud("SmallDataset",
-                                 download_directory=download_directory,
-                                 use_public_access=True)
+    dataset_cloud = DatasetCloud(
+        "SmallDataset", download_directory=download_directory, use_public_access=True
+    )
     dataset_list = dataset_cloud.get_existing_datasets()
     assert len(dataset_list) > 0, "Dataset list is empty."
-    assert "SmallDataset" in dataset_list,\
-        "Dataset list does not contain the dataset."
-        
-    # Test if the list does not contain duplicates
-    assert len(dataset_list) == len(set(dataset_list)),\
-        "Dataset list contains duplicates."
-
+    assert "SmallDataset" in dataset_list, "Dataset list does not contain the dataset."
 
+    # Test if the list does not contain duplicates
+    assert len(dataset_list) == len(
+        set(dataset_list)
+    ), "Dataset list contains duplicates."
 
 
 if "GOOGLE_APPLICATION_CREDENTIALS" in dict(environ):
+
     def test_update_dataset_list():
         # Create DatasetCloud object
         dataset_cloud = DatasetCloud("", use_public_access=False)
         # Update the dataset list
         dataset_cloud._update_dataset_list()
-        
-    
+
     def test_upload_and_download():
-        for data_format in ['pytorch_tensor', 'numpy_array']:
+        for data_format in ["pytorch_tensor", "numpy_array"]:
             download_directory = join("examples", "data", "DatasetCloud")
             # Generate a dataset
             # You don't have to do that if you already have a pickled dataset
             size_dataset = 100
             input_dim = 5
             num_labels = 2
-            
-            if data_format == 'pytorch_tensor':
+
+            if data_format == "pytorch_tensor":
                 data = torch.rand(size_dataset, input_dim)
                 labels = torch.randint(0, num_labels, (size_dataset,)).long()
 
                 # pickle data and labels
-                data_filename = 'tmp_data.pt'
-                labels_filename = 'tmp_labels.pt'
+                data_filename = "tmp_data.pt"
+                labels_filename = "tmp_labels.pt"
                 torch.save(data, data_filename)
                 torch.save(labels, labels_filename)
-            elif data_format == 'numpy_array':
+            elif data_format == "numpy_array":
                 data = np.random.rand(size_dataset, input_dim)
-                labels = np.random.randint(0, num_labels, (size_dataset,),
-                                           dtype=np.long)
+                labels = np.random.randint(
+                    0, num_labels, (size_dataset,), dtype=np.long
+                )
 
                 # pickle data and labels
-                data_filename = 'tmp_data.npy'
-                labels_filename = 'tmp_labels.npy'
+                data_filename = "tmp_data.npy"
+                labels_filename = "tmp_labels.npy"
                 np.save(data_filename, data)
                 np.save(labels_filename, labels)
             else:
@@ -111,9 +107,11 @@ def test_upload_and_download():
 
             ## Upload dataset to Cloud
             dataset_name = "TmpSmallDataset"
-            dataset_cloud = DatasetCloud(dataset_name,
-                                        download_directory=download_directory,
-                                        use_public_access=False)
+            dataset_cloud = DatasetCloud(
+                dataset_name,
+                download_directory=download_directory,
+                use_public_access=False,
+            )
 
             # Specify the metadata of the dataset
             dataset_cloud._add_metadata(
@@ -135,17 +133,17 @@ def test_upload_and_download():
             dataset_cloud._data_cloud.delete_blobs(dataset_name)
 
             # check whether downloaded dataset is the same as the original dataset
-            if data_format == 'pytorch_tensor':
+            if data_format == "pytorch_tensor":
                 downloaded_files = ["data.pt", "labels.pt", "metadata.json"]
-            elif data_format == 'numpy_array':
+            elif data_format == "numpy_array":
                 downloaded_files = ["data.npy", "labels.npy", "metadata.json"]
             for file in downloaded_files:
-                hash_original = get_checksum('tmp_' + file)
+                hash_original = get_checksum("tmp_" + file)
                 path_downloaded_file = join(download_directory, dataset_name, file)
                 hash_downloaded = get_checksum(path_downloaded_file)
-                assert hash_original == hash_downloaded, \
-                    "Original and downloaded files do not match."
-                
+                assert (
+                    hash_original == hash_downloaded
+                ), "Original and downloaded files do not match."
 
             # remove the labels and data files
             remove(data_filename)
diff --git a/gdeep/data/datasets/tests/test_dataset_for_image.py b/gdeep/data/datasets/tests/test_dataset_for_image.py
index f0039714..90a2b892 100644
--- a/gdeep/data/datasets/tests/test_dataset_for_image.py
+++ b/gdeep/data/datasets/tests/test_dataset_for_image.py
@@ -1,5 +1,8 @@
-from gdeep.data.datasets import ImageClassificationFromFiles, \
-    DataLoaderBuilder, DlBuilderFromDataCloud
+from gdeep.data.datasets import (
+    ImageClassificationFromFiles,
+    DataLoaderBuilder,
+    DlBuilderFromDataCloud,
+)
 
 from gdeep.data import TransformingDataset
 from gdeep.data.preprocessors import ToTensorImage
@@ -13,17 +16,17 @@
 LOGGER = logging.getLogger(__name__)
 
 
-
 def test_images_from_file():
     """test DatasetImageClassificationFromFiles"""
     file_path = os.path.dirname(os.path.realpath(__file__))
-    transform = ToTensorImage((32,32))  # this is already fitted
+    transform = ToTensorImage((32, 32))  # this is already fitted
     ds = ImageClassificationFromFiles(
-        os.path.join(file_path,"img_data"),
-        os.path.join(file_path,"img_data","labels.csv"))
+        os.path.join(file_path, "img_data"),
+        os.path.join(file_path, "img_data", "labels.csv"),
+    )
 
     tds = TransformingDataset(ds, transform)
-    dl, *_ = DataLoaderBuilder((tds,)).build(({"batch_size" : 2},))
+    dl, *_ = DataLoaderBuilder((tds,)).build(({"batch_size": 2},))
     assert len(next(iter(dl))[0].shape) == 4
 
 
@@ -31,11 +34,11 @@ def test_dlbuilderfromdatacloud():
     dataset_name = "SmallDataset"
     download_directory = join("examples", "data", "DatasetCloud")
 
-    dl_cloud_builder = DlBuilderFromDataCloud(dataset_name,
-                                    download_directory)
+    dl_cloud_builder = DlBuilderFromDataCloud(dataset_name, download_directory)
 
-    train_dataloader, val_dataloader, test_dataloader = \
-        dl_cloud_builder.build(({"batch_size" : 10},))
+    train_dataloader, val_dataloader, test_dataloader = dl_cloud_builder.build(
+        ({"batch_size": 10},)
+    )
     x, y = next(iter(train_dataloader))
     assert x.shape == torch.Size([10, 5])
-    assert y.shape == torch.Size([10])
\ No newline at end of file
+    assert y.shape == torch.Size([10])
diff --git a/gdeep/data/datasets/tests/test_dataset_from_array.py b/gdeep/data/datasets/tests/test_dataset_from_array.py
index 3ef619e7..3703e3d5 100644
--- a/gdeep/data/datasets/tests/test_dataset_from_array.py
+++ b/gdeep/data/datasets/tests/test_dataset_from_array.py
@@ -1,5 +1,9 @@
-from gdeep.data.datasets import ImageClassificationFromFiles, \
-    DataLoaderBuilder, FromArray, DlBuilderFromDataCloud
+from gdeep.data.datasets import (
+    ImageClassificationFromFiles,
+    DataLoaderBuilder,
+    FromArray,
+    DlBuilderFromDataCloud,
+)
 
 from gdeep.data.preprocessors import ToTensorImage
 import os
@@ -12,25 +16,24 @@
 LOGGER = logging.getLogger(__name__)
 
 
-
 def test_array():
     """test class DatasetFromArray"""
-    X = np.random.rand(10,4)
+    X = np.random.rand(10, 4)
     y = np.random.randint(3, size=10)
     ds = FromArray(X, y)
-    dl, *_ = DataLoaderBuilder((ds,)).build(({"batch_size" : 1},))
+    dl, *_ = DataLoaderBuilder((ds,)).build(({"batch_size": 1},))
     item = next(iter(dl))
     assert torch.norm(item[0] - torch.tensor(X[0])) < 1e-6
     assert item[1] == torch.tensor(y[0])
     assert len(dl.dataset) == 10
 
+
 def test_array_tensor():
     """test class DatasetFromArray"""
-    X = torch.rand(10,4)
+    X = torch.rand(10, 4)
     y = torch.randint(3, size=(10,))
     ds = FromArray(X, y)
     dl, *_ = DataLoaderBuilder((ds,)).build()
     item = next(iter(dl))
     assert torch.norm(item[0] - X[0]) < 1e-6
     assert item[1] == y[0]
-
diff --git a/gdeep/data/datasets/tests/test_parallel_orbit.py b/gdeep/data/datasets/tests/test_parallel_orbit.py
index faa0291a..aa23c38e 100644
--- a/gdeep/data/datasets/tests/test_parallel_orbit.py
+++ b/gdeep/data/datasets/tests/test_parallel_orbit.py
@@ -1,4 +1,3 @@
-
 import numpy as np
 
 from ..parallel_orbit import OrbitsGenerator
@@ -11,37 +10,46 @@
 og = OrbitsGenerator(
     parameters=parameters,
     num_orbits_per_class=num_orbits_per_class,
-    num_pts_per_orbit=num_pts_per_orbit
+    num_pts_per_orbit=num_pts_per_orbit,
 )
 
 orbits = og.get_orbits()
 
+
 def test_size_of_orbits():
-    """Tests if the size of the generator is correct
-    """
-    
-    
-    assert orbits.shape == (len(parameters) * num_orbits_per_class, num_pts_per_orbit, 2)
-    
+    """Tests if the size of the generator is correct"""
+
+    assert orbits.shape == (
+        len(parameters) * num_orbits_per_class,
+        num_pts_per_orbit,
+        2,
+    )
+
+
 def test_size_of_persistence_diagrams():
-    """Test if the size of the persistence diagram is correct
-    """
-    
+    """Test if the size of the persistence diagram is correct"""
+
     persistence_diagrams = og.get_persistence_diagrams()
-    
-    assert [persistence_diagrams.shape[i] for i in (0, 2)]\
-        == [len(parameters) * num_orbits_per_class, 2 + og._num_homology_dimensions]
+
+    assert [persistence_diagrams.shape[i] for i in (0, 2)] == [
+        len(parameters) * num_orbits_per_class,
+        2 + og._num_homology_dimensions,
+    ]
+
 
 def test_recursive_definition():
-    """Tests if the recursively defined orbits are correct
-    """
+    """Tests if the recursively defined orbits are correct"""
     for p_idx, p in enumerate(parameters):
         for o in range(num_orbits_per_class):
             for i in range(1, num_pts_per_orbit - 1):
-                x_current = orbits[p_idx * num_orbits_per_class + o, i    , 0]
-                x_next    = orbits[p_idx * num_orbits_per_class + o, i + 1, 0]
-                y_current = orbits[p_idx * num_orbits_per_class + o, i    , 1]
-                y_next    = orbits[p_idx * num_orbits_per_class + o, i + 1, 1]
-                
-                assert np.isclose(x_next, (x_current + p * y_current * (1.0 - y_current)) % 1.0)
-                assert np.isclose(y_next, (y_current + p * x_next * (1.0 - x_next)) % 1.0)
\ No newline at end of file
+                x_current = orbits[p_idx * num_orbits_per_class + o, i, 0]
+                x_next = orbits[p_idx * num_orbits_per_class + o, i + 1, 0]
+                y_current = orbits[p_idx * num_orbits_per_class + o, i, 1]
+                y_next = orbits[p_idx * num_orbits_per_class + o, i + 1, 1]
+
+                assert np.isclose(
+                    x_next, (x_current + p * y_current * (1.0 - y_current)) % 1.0
+                )
+                assert np.isclose(
+                    y_next, (y_current + p * x_next * (1.0 - x_next)) % 1.0
+                )
diff --git a/gdeep/data/datasets/tests/test_persistence_diagrams_from_graphs_builder.py b/gdeep/data/datasets/tests/test_persistence_diagrams_from_graphs_builder.py
index fc97f9f2..02e7b4e7 100644
--- a/gdeep/data/datasets/tests/test_persistence_diagrams_from_graphs_builder.py
+++ b/gdeep/data/datasets/tests/test_persistence_diagrams_from_graphs_builder.py
@@ -1,44 +1,61 @@
 import os
 
 import numpy as np
-from gdeep.data.datasets.persistence_diagrams_from_graphs_builder import PersistenceDiagramFromGraphBuilder
-from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import OneHotEncodedPersistenceDiagram,\
-    get_one_hot_encoded_persistence_diagram_from_gudhi_extended
+from gdeep.data.datasets.persistence_diagrams_from_graphs_builder import (
+    PersistenceDiagramFromGraphBuilder,
+)
+from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import (
+    OneHotEncodedPersistenceDiagram,
+    get_one_hot_encoded_persistence_diagram_from_gudhi_extended,
+)
 from gdeep.utility import DEFAULT_GRAPH_DIR
 
 
-
 def test_mutag_load_and_save():
 
     # Parameters
-    name_graph_dataset: str = 'MUTAG'
+    name_graph_dataset: str = "MUTAG"
     diffusion_parameter: float = 0.1
     num_homology_types: int = 4
 
-
     # Create the persistence diagram dataset
-    pd_creator = PersistenceDiagramFromGraphBuilder(name_graph_dataset, diffusion_parameter)
+    pd_creator = PersistenceDiagramFromGraphBuilder(
+        name_graph_dataset, diffusion_parameter
+    )
     pd_creator.create()
 
-    gudhi_ext = \
-    (np.array([[0.90634646, 0.90835863],
-            [0.90659781, 0.90835863],
-            [0.90672356, 0.90684934],
-            [0.90672356, 0.90684934]]),
-    np.array([[0.90634646, 0.90835863]]),
-    np.array([[0.90659795, 0.90710085],
-            [0.90672356, 0.90710077],
-            [0.90672356, 0.90684934]]),
-    np.array([[0.90659781, 0.90710077],
-            [0.90659781, 0.90710085]]))
+    gudhi_ext = (
+        np.array(
+            [
+                [0.90634646, 0.90835863],
+                [0.90659781, 0.90835863],
+                [0.90672356, 0.90684934],
+                [0.90672356, 0.90684934],
+            ]
+        ),
+        np.array([[0.90634646, 0.90835863]]),
+        np.array(
+            [
+                [0.90659795, 0.90710085],
+                [0.90672356, 0.90710077],
+                [0.90672356, 0.90684934],
+            ]
+        ),
+        np.array([[0.90659781, 0.90710077], [0.90659781, 0.90710085]]),
+    )
 
     pd_gudhi = get_one_hot_encoded_persistence_diagram_from_gudhi_extended(gudhi_ext)
 
     # Print sample extended persistence diagram
-    file_path: str = os.path.join(DEFAULT_GRAPH_DIR,
-                                f"MUTAG_{diffusion_parameter}_extended_persistence", "diagrams")
+    file_path: str = os.path.join(
+        DEFAULT_GRAPH_DIR,
+        f"MUTAG_{diffusion_parameter}_extended_persistence",
+        "diagrams",
+    )
     graph_idx = 1
-    pd = OneHotEncodedPersistenceDiagram.load(os.path.join(file_path, 
-                                                        f"{graph_idx}.npy"))
-    assert pd_gudhi.all_close(pd, atol=1e-6), \
-        "Generated persistence diagram is not equal to the one from GUDHI"
\ No newline at end of file
+    pd = OneHotEncodedPersistenceDiagram.load(
+        os.path.join(file_path, f"{graph_idx}.npy")
+    )
+    assert pd_gudhi.all_close(
+        pd, atol=1e-6
+    ), "Generated persistence diagram is not equal to the one from GUDHI"
diff --git a/gdeep/data/persistence_diagrams/tests/test_one_hot_persistence_diagram.py b/gdeep/data/persistence_diagrams/tests/test_one_hot_persistence_diagram.py
index 4649b3e6..0e4c8ca6 100644
--- a/gdeep/data/persistence_diagrams/tests/test_one_hot_persistence_diagram.py
+++ b/gdeep/data/persistence_diagrams/tests/test_one_hot_persistence_diagram.py
@@ -2,17 +2,21 @@
 import torch
 from gtda.homology import VietorisRipsPersistence
 
-from .. import (OneHotEncodedPersistenceDiagram,
-                get_one_hot_encoded_persistence_diagram_from_gtda)
+from .. import (
+    OneHotEncodedPersistenceDiagram,
+    get_one_hot_encoded_persistence_diagram_from_gtda,
+)
 
 
 def test_one_hot_encoded_persistence_diagram():
     pd_one_hot = OneHotEncodedPersistenceDiagram(
-        torch.tensor([
-            [0.3, 0.5, 1.0, 0.0],
-            [0.4, 0.8, 1.0, 0.0],
-            [0.5, 0.9, 0.0, 1.0],
-        ])
+        torch.tensor(
+            [
+                [0.3, 0.5, 1.0, 0.0],
+                [0.4, 0.8, 1.0, 0.0],
+                [0.5, 0.9, 0.0, 1.0],
+            ]
+        )
     )
 
     lifetimes = pd_one_hot.get_lifetimes()
@@ -21,25 +25,38 @@ def test_one_hot_encoded_persistence_diagram():
 
     pd_filtered = pd_one_hot.filter_by_lifetime(min_lifetime=0.3, max_lifetime=0.5)
 
-    assert pd_filtered.all_close(OneHotEncodedPersistenceDiagram(
-        torch.tensor([[0.5000, 0.9000, 0.0000, 1.0000],
-            [0.4000, 0.8000, 1.0000, 0.0000]]))
+    assert pd_filtered.all_close(
+        OneHotEncodedPersistenceDiagram(
+            torch.tensor(
+                [[0.5000, 0.9000, 0.0000, 1.0000], [0.4000, 0.8000, 1.0000, 0.0000]]
+            )
+        )
     )
 
+
 def test_from_numpy():
-    x = np.array([
+    x = np.array(
+        [
             [0.3, 0.5, 1.0, 0.0],
             [0.4, 0.8, 1.0, 0.0],
             [0.5, 0.9, 0.0, 1.0],
-        ])
+        ]
+    )
 
     x = OneHotEncodedPersistenceDiagram.from_numpy(x)
 
-    assert x.all_close(OneHotEncodedPersistenceDiagram(
-                        torch.tensor([[0.3000, 0.5000, 1.0000, 0.0000],
-                                      [0.5000, 0.9000, 0.0000, 1.0000],
-                                      [0.4000, 0.8000, 1.0000, 0.0000]])
-    ))
+    assert x.all_close(
+        OneHotEncodedPersistenceDiagram(
+            torch.tensor(
+                [
+                    [0.3000, 0.5000, 1.0000, 0.0000],
+                    [0.5000, 0.9000, 0.0000, 1.0000],
+                    [0.4000, 0.8000, 1.0000, 0.0000],
+                ]
+            )
+        )
+    )
+
 
 def test_get_one_hot_encoded_persistence_diagram_from_gtda():
     vr = VietorisRipsPersistence(homology_dimensions=[0, 1])
diff --git a/gdeep/data/preprocessors/tests/__init__.py b/gdeep/data/preprocessors/tests/__init__.py
index 24fa198e..3d57db5a 100644
--- a/gdeep/data/preprocessors/tests/__init__.py
+++ b/gdeep/data/preprocessors/tests/__init__.py
@@ -1,1 +1,1 @@
-#from ..data_cloud import DataCloud
+# from ..data_cloud import DataCloud
diff --git a/gdeep/data/preprocessors/tests/test_preprocessing_pipeline.py b/gdeep/data/preprocessors/tests/test_preprocessing_pipeline.py
index 3faf63a5..46d7e362 100644
--- a/gdeep/data/preprocessors/tests/test_preprocessing_pipeline.py
+++ b/gdeep/data/preprocessors/tests/test_preprocessing_pipeline.py
@@ -4,19 +4,25 @@
 import torch
 
 from gdeep.data.datasets import ImageClassificationFromFiles
-from gdeep.data.datasets.persistence_diagrams_from_files import \
-    PersistenceDiagramFromFiles
-from gdeep.data.datasets.persistence_diagrams_from_graphs_builder import \
-    PersistenceDiagramFromGraphBuilder
-from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import \
-    OneHotEncodedPersistenceDiagram
+from gdeep.data.datasets.persistence_diagrams_from_files import (
+    PersistenceDiagramFromFiles,
+)
+from gdeep.data.datasets.persistence_diagrams_from_graphs_builder import (
+    PersistenceDiagramFromGraphBuilder,
+)
+from gdeep.data.persistence_diagrams.one_hot_persistence_diagram import (
+    OneHotEncodedPersistenceDiagram,
+)
 from gdeep.data.preprocessors import Normalization, ToTensorImage
-from gdeep.data.preprocessors.filter_persistence_diagram_by_homology_dimension import \
-    FilterPersistenceDiagramByHomologyDimension
-from gdeep.data.preprocessors.filter_persistence_diagram_by_lifetime import \
-    FilterPersistenceDiagramByLifetime
-from gdeep.data.preprocessors.normalization_persistence_diagram import \
-    NormalizationPersistenceDiagram
+from gdeep.data.preprocessors.filter_persistence_diagram_by_homology_dimension import (
+    FilterPersistenceDiagramByHomologyDimension,
+)
+from gdeep.data.preprocessors.filter_persistence_diagram_by_lifetime import (
+    FilterPersistenceDiagramByLifetime,
+)
+from gdeep.data.preprocessors.normalization_persistence_diagram import (
+    NormalizationPersistenceDiagram,
+)
 from gdeep.utility.constants import DEFAULT_GRAPH_DIR
 from gdeep.data.preprocessors.normalization import _compute_mean_of_dataset
 from sklearn import preprocessing
@@ -30,38 +36,42 @@
 def test_preeprocessing_pipeline():
     """test preprocessing pipeline"""
     normalize = Normalization()
-    to_tensor = ToTensorImage([32,32])
+    to_tensor = ToTensorImage([32, 32])
     file_path = os.path.dirname(os.path.realpath(__file__))
 
     image_dataset = ImageClassificationFromFiles(
         os.path.join(file_path, "img_data"),
-        os.path.join(file_path, "img_data", "labels.csv"))
+        os.path.join(file_path, "img_data", "labels.csv"),
+    )
     # define preprocessing pipeline
     preprocessing_pipeline = PreprocessingPipeline((to_tensor, normalize))
 
     # fit preprocessing pipeline to dataset
     preprocessing_pipeline.fit_to_dataset(image_dataset)
-    
+
     # apply preprocessing pipeline to dataset
-    preprocessed_dataset = preprocessing_pipeline.attach_transform_to_dataset(image_dataset)
-    
+    preprocessed_dataset = preprocessing_pipeline.attach_transform_to_dataset(
+        image_dataset
+    )
+
     assert len(preprocessed_dataset[0][0].shape) == 3
-    
+
 
 def test_preprocessing_pipeline_for_persistence_diagrams():
     # Parameters
-    name_graph_dataset: str = 'MUTAG'
+    name_graph_dataset: str = "MUTAG"
     diffusion_parameter: float = 0.1
     num_homology_types: int = 4
 
-
     # Create the persistence diagram dataset
-    pd_creator = PersistenceDiagramFromGraphBuilder(name_graph_dataset, diffusion_parameter)
+    pd_creator = PersistenceDiagramFromGraphBuilder(
+        name_graph_dataset, diffusion_parameter
+    )
     pd_creator.create()
-    
+
     pd_mutag_ds = PersistenceDiagramFromFiles(
-    os.path.join(
-        DEFAULT_GRAPH_DIR, f"MUTAG_{diffusion_parameter}_extended_persistence"
+        os.path.join(
+            DEFAULT_GRAPH_DIR, f"MUTAG_{diffusion_parameter}_extended_persistence"
         )
     )
     # Create the train/validation/test split
@@ -72,7 +82,7 @@ def test_preprocessing_pipeline_for_persistence_diagrams():
         random_state=42,
     )
 
-    train_indices , validation_indices = train_test_split(
+    train_indices, validation_indices = train_test_split(
         train_indices,
         test_size=0.2,
         random_state=42,
@@ -84,28 +94,37 @@ def test_preprocessing_pipeline_for_persistence_diagrams():
     test_dataset = Subset(pd_mutag_ds, test_indices)
 
     # Preprocess the data
-    preprocessing_pipeline = PreprocessingPipeline[Tuple[OneHotEncodedPersistenceDiagram, int]](
+    preprocessing_pipeline = PreprocessingPipeline[
+        Tuple[OneHotEncodedPersistenceDiagram, int]
+    ](
         (
             FilterPersistenceDiagramByHomologyDimension[int]([0, 1]),
-            FilterPersistenceDiagramByLifetime[int](min_lifetime=-0.1, max_lifetime=1.0),
+            FilterPersistenceDiagramByLifetime[int](
+                min_lifetime=-0.1, max_lifetime=1.0
+            ),
             NormalizationPersistenceDiagram[int](num_homology_dimensions=4),
         )
     )
 
     preprocessing_pipeline.fit_to_dataset(train_dataset)
-    
+
     train_dataset = preprocessing_pipeline.attach_transform_to_dataset(train_dataset)
-    validation_dataset = preprocessing_pipeline.attach_transform_to_dataset(validation_dataset)
+    validation_dataset = preprocessing_pipeline.attach_transform_to_dataset(
+        validation_dataset
+    )
     test_dataset = preprocessing_pipeline.attach_transform_to_dataset(test_dataset)
 
-    
     computed_mean = _compute_mean_of_dataset(
-                TransformingDataset(train_dataset, lambda x: (x[0].get_raw_data().mean(dim=0), x[1]))
-                )[:2]
+        TransformingDataset(
+            train_dataset, lambda x: (x[0].get_raw_data().mean(dim=0), x[1])
+        )
+    )[:2]
     # -> (0.0, 0.0, sth, sth)
     assert torch.allclose(computed_mean.float(), torch.tensor([0.0, 0.0]), atol=1e-4)
-    computed_stddev= _compute_mean_of_dataset(
-                TransformingDataset(train_dataset, lambda x: ((x[0].get_raw_data()**2).mean(dim=0), x[1]))
-                )[:2]
+    computed_stddev = _compute_mean_of_dataset(
+        TransformingDataset(
+            train_dataset, lambda x: ((x[0].get_raw_data() ** 2).mean(dim=0), x[1])
+        )
+    )[:2]
     # -> (1.0, 1.0, sth, sth)
     assert torch.allclose(computed_stddev.float(), torch.tensor([1.0, 1.0]), atol=1e-4)
diff --git a/gdeep/models/tests/test_periodic_nn.py b/gdeep/models/tests/test_periodic_nn.py
index d25565b5..ab45b0ce 100644
--- a/gdeep/models/tests/test_periodic_nn.py
+++ b/gdeep/models/tests/test_periodic_nn.py
@@ -8,7 +8,7 @@
 def test_periodic():
     X = torch.rand(10, 2)
     sn = FFNet([2, 3, 2])
-    bdry_list = [(-1,1),(0.5,1)]
+    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/search/tests/data/SmallMutagDataset/metadata.json b/gdeep/search/tests/data/SmallMutagDataset/metadata.json
index 13e135df..4654a779 100644
--- a/gdeep/search/tests/data/SmallMutagDataset/metadata.json
+++ b/gdeep/search/tests/data/SmallMutagDataset/metadata.json
@@ -10,4 +10,4 @@
     "num_labels": 2,
     "size": 5,
     "task_type": "classification"
-}
\ No newline at end of file
+}
diff --git a/gdeep/search/tests/test_benchmark.py b/gdeep/search/tests/test_benchmark.py
index e67af739..df99076c 100644
--- a/gdeep/search/tests/test_benchmark.py
+++ b/gdeep/search/tests/test_benchmark.py
@@ -15,9 +15,12 @@
 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))
+        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)
@@ -63,12 +66,16 @@ def forward(self, x):
 dataloaders_dicts = []
 
 
-test_indices = [64*5 + x for x in range(32*3)]
-train_indices = [x for x in range(32*2)]
+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)},  # type: ignore
-                             {"batch_size": 32, "sampler": SubsetRandomSampler(test_indices)}])
+dl_tr, dl_val, _ = dl.build(
+    [
+        {"batch_size": 32, "sampler": SubsetRandomSampler(train_indices)},  # type: ignore
+        {"batch_size": 32, "sampler": SubsetRandomSampler(test_indices)},
+    ]
+)
 
 
 temp_dict = {"name": "CIFAR10_1000", "dataloaders": [dl_tr, dl_val]}  # type: ignore
diff --git a/gdeep/search/tests/test_hpo.py b/gdeep/search/tests/test_hpo.py
index 4a3b3bed..c444f062 100644
--- a/gdeep/search/tests/test_hpo.py
+++ b/gdeep/search/tests/test_hpo.py
@@ -22,9 +22,12 @@
 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))
+        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)
@@ -32,7 +35,6 @@ def forward(self, x):
 
 # model
 class Model1(nn.Module):
-
     def __init__(self):
         super(Model1, self).__init__()
         self.seqmodel = nn.Sequential(nn.Flatten(), FFNet(arch=[4, 5, 4]))
@@ -56,7 +58,8 @@ def forward(self, x):
 # use only 32*7 images from CIFAR10
 train_indices = list(range(32 * 7))
 dl_tr, *_ = DataLoaderBuilder([transformed_ds_tr]).build(  # type: ignore
-    [{"batch_size": 32, "sampler": SubsetRandomSampler(train_indices)}])
+    [{"batch_size": 32, "sampler": SubsetRandomSampler(train_indices)}]
+)
 
 
 def test_hpo_failure():
@@ -67,8 +70,13 @@ def test_hpo_failure():
     loss_fn = nn.CrossEntropyLoss()
 
     # initialise pipeline class
-    pipe = Trainer(model, [dl_tr, None], loss_fn, writer,
-                   k_fold_class=StratifiedKFold(2, shuffle=True))
+    pipe = Trainer(
+        model,
+        [dl_tr, None],
+        loss_fn,
+        writer,
+        k_fold_class=StratifiedKFold(2, shuffle=True),
+    )
 
     # initialise gridsearch
     try:
@@ -85,8 +93,13 @@ def test_hpo_cross_val():
     loss_fn = nn.CrossEntropyLoss()
 
     # initialise pipeline class
-    pipe = Trainer(model, [dl_tr, None], loss_fn, writer,
-                   k_fold_class=StratifiedKFold(2, shuffle=True))
+    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)
@@ -100,8 +113,15 @@ def test_hpo_cross_val():
     models_hyperparams = {"n_nodes": ["200"]}
 
     # starting the gridsearch
-    search.start([SGD, Adam], 3, False, optimizers_params, dataloaders_params,
-                 models_hyperparams, n_accumulated_grads=2)
+    search.start(
+        [SGD, Adam],
+        3,
+        False,
+        optimizers_params,
+        dataloaders_params,
+        models_hyperparams,
+        n_accumulated_grads=2,
+    )
 
 
 def test_hpo_accumulated_grads():
@@ -126,8 +146,15 @@ def test_hpo_accumulated_grads():
     models_hyperparams = {"n_nodes": ["200"]}
 
     # starting the gridsearch
-    search.start([SGD, Adam], 2, False, optimizers_params, dataloaders_params,
-                 models_hyperparams, n_accumulated_grads=2)
+    search.start(
+        [SGD, Adam],
+        2,
+        False,
+        optimizers_params,
+        dataloaders_params,
+        models_hyperparams,
+        n_accumulated_grads=2,
+    )
 
 
 def test_hpo_loss():
@@ -149,7 +176,9 @@ def test_hpo_loss():
     models_hyperparams = {"n_nodes": ["200"]}
 
     # starting the gridsearch
-    search.start([SGD, Adam], 2, False, optimizers_params, dataloaders_params, models_hyperparams)
+    search.start(
+        [SGD, Adam], 2, False, optimizers_params, dataloaders_params, models_hyperparams
+    )
 
 
 def test_hpo_string_parameters():
@@ -174,7 +203,9 @@ def test_hpo_string_parameters():
     models_hyperparams = {"n_nodes": ["200", "256"]}
 
     # starting the gridsearch
-    search.start([SGD], 1, False, optimizers_params, dataloaders_params, models_hyperparams)
+    search.start(
+        [SGD], 1, False, optimizers_params, dataloaders_params, models_hyperparams
+    )
 
 
 def test_hpo_collate():
@@ -183,7 +214,9 @@ def __init__(self):
             self.x = []
             for _ in range(100):
                 self.x.append((torch.rand(1, np.random.randint(2, 4))).to(torch.float))
-            self.y = np.array(np.random.randint(2, size=100 * 2).reshape(-1, 2), dtype=np.int64)
+            self.y = np.array(
+                np.random.randint(2, size=100 * 2).reshape(-1, 2), dtype=np.int64
+            )
 
         def __len__(self):
             return 100
@@ -197,21 +230,19 @@ def collate_fn(batch_tuple: List):
         label = torch.zeros(len(batch_tuple), 2).to(torch.long)
         for i, batch in enumerate(batch_tuple):
             source = batch[0]
-            target[i, :len(source[-1])] = source
+            target[i, : len(source[-1])] = source
             label[i] = torch.tensor(batch[1]).to(torch.long)
         return target, label
 
     model = Model1()
     # dataloaders
     ds = MyDataset()
-    dl_train = DataLoader(ds, batch_size=6,
-                       collate_fn=collate_fn)
+    dl_train = DataLoader(ds, batch_size=6, collate_fn=collate_fn)
 
     # loss function
     loss_fn = nn.CrossEntropyLoss()
     # pipeline
-    pipe = Trainer(model, [dl_train, None],
-                   loss_fn, writer)
+    pipe = Trainer(model, [dl_train, None], loss_fn, writer)
     # initialise gridsearch
     search = HyperParameterOptimization(pipe, "loss", 2, best_not_last=True)
 
diff --git a/gdeep/search/tests/test_hpo_config.py b/gdeep/search/tests/test_hpo_config.py
index 7af369df..34930754 100644
--- a/gdeep/search/tests/test_hpo_config.py
+++ b/gdeep/search/tests/test_hpo_config.py
@@ -1,4 +1,3 @@
-
 from gdeep.search import HPOConfig
 from torch.optim import SGD
 
diff --git a/gdeep/topology_layers/tests/test_persformer.py b/gdeep/topology_layers/tests/test_persformer.py
index a27a7c69..dfa4cf25 100644
--- a/gdeep/topology_layers/tests/test_persformer.py
+++ b/gdeep/topology_layers/tests/test_persformer.py
@@ -4,15 +4,22 @@
 import torch.nn as nn
 from gdeep.data import PreprocessingPipeline
 from gdeep.data.datasets import PersistenceDiagramFromFiles
-from gdeep.data.datasets.base_dataloaders import (DataLoaderBuilder,
-                                                  DataLoaderParamsTuples)
-from gdeep.data.datasets.persistence_diagrams_from_graphs_builder import \
-    PersistenceDiagramFromGraphBuilder
+from gdeep.data.datasets.base_dataloaders import (
+    DataLoaderBuilder,
+    DataLoaderParamsTuples,
+)
+from gdeep.data.datasets.persistence_diagrams_from_graphs_builder import (
+    PersistenceDiagramFromGraphBuilder,
+)
 from gdeep.data.persistence_diagrams import (
-    OneHotEncodedPersistenceDiagram, collate_fn_persistence_diagrams)
+    OneHotEncodedPersistenceDiagram,
+    collate_fn_persistence_diagrams,
+)
 from gdeep.data.preprocessors import (
     FilterPersistenceDiagramByHomologyDimension,
-    FilterPersistenceDiagramByLifetime, NormalizationPersistenceDiagram)
+    FilterPersistenceDiagramByLifetime,
+    NormalizationPersistenceDiagram,
+)
 from gdeep.topology_layers import Persformer, PersformerConfig
 from gdeep.trainer.trainer import Trainer
 from gdeep.utility import DEFAULT_GRAPH_DIR, PoolerType
@@ -26,12 +33,14 @@
 
 def test_persformer_training():
     # Parameters
-    name_graph_dataset: str = 'MUTAG'
+    name_graph_dataset: str = "MUTAG"
     diffusion_parameter: float = 0.1
     num_homology_types: int = 4
 
     # Create the persistence diagram dataset
-    pd_creator = PersistenceDiagramFromGraphBuilder(name_graph_dataset, diffusion_parameter)
+    pd_creator = PersistenceDiagramFromGraphBuilder(
+        name_graph_dataset, diffusion_parameter
+    )
     pd_creator.create()
 
     pd_mutag_ds = PersistenceDiagramFromFiles(
@@ -59,10 +68,14 @@ def test_persformer_training():
     test_dataset = Subset(pd_mutag_ds, test_indices)
 
     # Preprocess the data
-    preprocessing_pipeline = PreprocessingPipeline[Tuple[OneHotEncodedPersistenceDiagram, int]](
+    preprocessing_pipeline = PreprocessingPipeline[
+        Tuple[OneHotEncodedPersistenceDiagram, int]
+    ](
         (
             FilterPersistenceDiagramByHomologyDimension[int]([0, 1]),
-            FilterPersistenceDiagramByLifetime[int](min_lifetime=-0.1, max_lifetime=1.0),
+            FilterPersistenceDiagramByLifetime[int](
+                min_lifetime=-0.1, max_lifetime=1.0
+            ),
             NormalizationPersistenceDiagram[int](num_homology_dimensions=4),
         )
     )
@@ -70,7 +83,9 @@ def test_persformer_training():
     preprocessing_pipeline.fit_to_dataset(train_dataset)
 
     train_dataset = preprocessing_pipeline.attach_transform_to_dataset(train_dataset)
-    validation_dataset = preprocessing_pipeline.attach_transform_to_dataset(validation_dataset)
+    validation_dataset = preprocessing_pipeline.attach_transform_to_dataset(
+        validation_dataset
+    )
     test_dataset = preprocessing_pipeline.attach_transform_to_dataset(test_dataset)
 
     dl_params = DataLoaderParamsTuples.default(
@@ -100,6 +115,10 @@ def test_persformer_training():
 
     trainer = Trainer(model, [dl_train, dl_val, dl_test], loss_function, writer)
 
-    trainer.train(Adam, 3, False,
-                  {"lr": 0.01},
-                  {"batch_size": 16, "collate_fn": collate_fn_persistence_diagrams})
+    trainer.train(
+        Adam,
+        3,
+        False,
+        {"lr": 0.01},
+        {"batch_size": 16, "collate_fn": collate_fn_persistence_diagrams},
+    )
diff --git a/gdeep/trainer/tests/test_trainer.py b/gdeep/trainer/tests/test_trainer.py
index 4893dbd3..7c8a34dc 100644
--- a/gdeep/trainer/tests/test_trainer.py
+++ b/gdeep/trainer/tests/test_trainer.py
@@ -20,7 +20,9 @@ def __init__(self):
         self.x = []
         for _ in range(100):
             self.x.append((torch.rand(1, np.random.randint(2, 4))).to(torch.float))
-        self.y = np.array(np.random.randint(2, size=100 * 2).reshape(-1, 2), dtype=np.int64)
+        self.y = np.array(
+            np.random.randint(2, size=100 * 2).reshape(-1, 2), dtype=np.int64
+        )
 
     def __len__(self):
         return 100
@@ -35,14 +37,13 @@ def collate_fn(batch_tuple: List):
     label = torch.zeros(len(batch_tuple), 2).to(torch.long)
     for i, batch in enumerate(batch_tuple):
         source = batch[0]
-        target[i, :len(source[-1])] = source
+        target[i, : len(source[-1])] = source
         label[i] = torch.tensor(batch[1]).to(torch.long)
     return target, label
 
 
 # model
 class Model1(nn.Module):
-
     def __init__(self):
         super(Model1, self).__init__()
         self.seqmodel = nn.Sequential(nn.Flatten(), FFNet(arch=[4, 5, 4]))
@@ -57,7 +58,7 @@ def test_trainer_from_array():
     model = Model1()
     # dataloaders
     X = np.array(np.random.rand(100, 4), dtype=np.float32)
-    y = np.array(np.random.randint(2, size=100*2).reshape(-1, 2), dtype=np.int64)
+    y = np.array(np.random.randint(2, size=100 * 2).reshape(-1, 2), dtype=np.int64)
     dl_tr, *_ = DataLoaderBuilder([FromArray(X, y)]).build([{"batch_size": 23}])
 
     # loss function
@@ -65,8 +66,7 @@ def test_trainer_from_array():
     # tb writer
     writer = GiottoSummaryWriter()
     # pipeline
-    pipe = Trainer(model, [dl_tr, None],
-                   loss_fn, writer)
+    pipe = Trainer(model, [dl_tr, None], 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)
 
@@ -76,16 +76,14 @@ def test_trainer_collate():
     model = Model1()
     # dataloaders
     ds = MyDataset()
-    dl_tr = DataLoader(ds, batch_size=6,
-                       collate_fn=collate_fn)
+    dl_tr = DataLoader(ds, batch_size=6, collate_fn=collate_fn)
 
     # loss function
     loss_fn = nn.CrossEntropyLoss()
     # tb writer
     writer = GiottoSummaryWriter()
     # pipeline
-    pipe = Trainer(model, [dl_tr, None],
-                   loss_fn, writer)
+    pipe = Trainer(model, [dl_tr, None], 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
index 7c7aa5d8..b8a48ecf 100644
--- a/gdeep/trainer/tests/test_trainer_config.py
+++ b/gdeep/trainer/tests/test_trainer_config.py
@@ -1,4 +1,3 @@
-
 from gdeep.trainer import TrainerConfig
 from torch.optim import SGD
 
diff --git a/gdeep/utility/optimisation/tests/__init__.py b/gdeep/utility/optimisation/tests/__init__.py
index 8b137891..e69de29b 100644
--- a/gdeep/utility/optimisation/tests/__init__.py
+++ b/gdeep/utility/optimisation/tests/__init__.py
@@ -1,1 +0,0 @@
-
diff --git a/gdeep/utility/optimisation/tests/test_persistence_grad.py b/gdeep/utility/optimisation/tests/test_persistence_grad.py
index 72a99939..1eb8cb97 100644
--- a/gdeep/utility/optimisation/tests/test_persistence_grad.py
+++ b/gdeep/utility/optimisation/tests/test_persistence_grad.py
@@ -1,17 +1,14 @@
-
 import torch
 from ...optimisation import PersistenceGradient
 
 
 def test_persistence_gradient_2d():
-    """check if the class is consistent.
-    """
-    X = torch.tensor([[1, 0.], [0, 1.], [2, 2], [2, 1]])
+    """check if the class is consistent."""
+    X = torch.tensor([[1, 0.0], [0, 1.0], [2, 2], [2, 1]])
     hom_dim = [0, 1]
-    pg = PersistenceGradient(homology_dimensions=hom_dim,
-                             zeta=0.1,
-                             max_edge_length=3,
-                             collapse_edges=False)
+    pg = PersistenceGradient(
+        homology_dimensions=hom_dim, zeta=0.1, max_edge_length=3, collapse_edges=False
+    )
 
     assert X.dtype == torch.float32
     assert pg.phi(X).shape[0] == 14
@@ -21,14 +18,12 @@ def test_persistence_gradient_2d():
 
 
 def test_persistence_gradient_3d():
-    """check if the class is consistent.
-    """
-    X = torch.tensor([[1, 0., 1], [0, 1., 0], [2, 2, 1], [2, 1, 2]])
+    """check if the class is consistent."""
+    X = torch.tensor([[1, 0.0, 1], [0, 1.0, 0], [2, 2, 1], [2, 1, 2]])
     hom_dim = [0, 1]
-    pg = PersistenceGradient(homology_dimensions=hom_dim,
-                             zeta=0.1,
-                             max_edge_length=3,
-                             collapse_edges=False)
+    pg = PersistenceGradient(
+        homology_dimensions=hom_dim, zeta=0.1, max_edge_length=3, collapse_edges=False
+    )
     assert X.dtype == torch.float32
     assert pg.phi(X).shape[0] == 14
 
@@ -37,15 +32,14 @@ def test_persistence_gradient_3d():
 
 
 def test_persistence_gradient_5d():
-    """check if the class is consistent.
-    """
-    X = torch.tensor([[1, 0., 1, 0.5, 1], [0, 1., 0, 0.5, 1],
-                      [2, 2, 1, 0.5, 1], [2, 1, 2, 0.5, 1]])
+    """check if the class is consistent."""
+    X = torch.tensor(
+        [[1, 0.0, 1, 0.5, 1], [0, 1.0, 0, 0.5, 1], [2, 2, 1, 0.5, 1], [2, 1, 2, 0.5, 1]]
+    )
     hom_dim = [0, 1]
-    pg = PersistenceGradient(homology_dimensions=hom_dim,
-                             zeta=0.1,
-                             max_edge_length=3,
-                             collapse_edges=False)
+    pg = PersistenceGradient(
+        homology_dimensions=hom_dim, zeta=0.1, max_edge_length=3, collapse_edges=False
+    )
     assert X.dtype == torch.float32
     assert pg.phi(X).shape[0] == 14
     assert pg.persistence_function(X).item() >= -2.7783148288726807 - 0.001
@@ -54,14 +48,17 @@ def test_persistence_gradient_5d():
 
 def test_persistence_gradient_4d():
     """check if the matrix input works properly"""
-    X = torch.tensor([[1, 0., 1, 0.5], [0, 1., 0, 0.5],
-                      [2, 2, 1, 0.5], [2, 1, 2, 0.5]])
+    X = torch.tensor(
+        [[1, 0.0, 1, 0.5], [0, 1.0, 0, 0.5], [2, 2, 1, 0.5], [2, 1, 2, 0.5]]
+    )
     hom_dim = [0, 1]
-    pg = PersistenceGradient(homology_dimensions=hom_dim,
-                             zeta=0.1,
-                             max_edge_length=3,
-                             metric="precomputed",
-                             collapse_edges=False)
+    pg = PersistenceGradient(
+        homology_dimensions=hom_dim,
+        zeta=0.1,
+        max_edge_length=3,
+        metric="precomputed",
+        collapse_edges=False,
+    )
     assert X.dtype == torch.float32
     assert pg.phi(X).shape[0] == 14
     assert pg.persistence_function(X).item() >= 0.3467579483985901 + 0.001
@@ -70,51 +67,50 @@ def test_persistence_gradient_4d():
 
 def test_persistence_gradient_matrix():
     # simulate the weighted graph
-    dist = torch.tensor([[0., 2, 3],
-                        [2, 0., 2.2],
-                        [3, 2.2, 0.]])
-
-    pg = PersistenceGradient(homology_dimensions=[0, 1],
-                             zeta=0.0,
-                             collapse_edges=False,
-                             metric="precomputed")
-    assert all(pg.phi(dist) == torch.tensor([0., 0., 0., 2., 2.2, 3., 3.]))
+    dist = torch.tensor([[0.0, 2, 3], [2, 0.0, 2.2], [3, 2.2, 0.0]])
+
+    pg = PersistenceGradient(
+        homology_dimensions=[0, 1], zeta=0.0, collapse_edges=False, metric="precomputed"
+    )
+    assert all(pg.phi(dist) == torch.tensor([0.0, 0.0, 0.0, 2.0, 2.2, 3.0, 3.0]))
     assert pg.persistence_function(dist).item() >= -6.3 - 0.0001
     pg.sgd(dist, n_epochs=1, lr=1)
 
 
 def test_persistence_gradient_matrix_2():
     # simulate the weighted graph
-    dist = torch.tensor([[0., 2., 10., 10.],
-                         [2., 0., 2., 1],
-                         [10., 2., 0., 1],
-                         [10., 1, 1, 0.]])
-    pg = PersistenceGradient(homology_dimensions=[0, 1],
-                             zeta=0.0,
-                             collapse_edges=False,
-                             metric="precomputed")
-    assert all(pg.phi(dist) == torch.tensor([0., 0., 0., 0., 1.,
-                                             1., 2., 2., 2., 10.,
-                                             10., 10., 10., 10.]))
-    assert pg.persistence_function(dist).item() > -23.
+    dist = torch.tensor(
+        [
+            [0.0, 2.0, 10.0, 10.0],
+            [2.0, 0.0, 2.0, 1],
+            [10.0, 2.0, 0.0, 1],
+            [10.0, 1, 1, 0.0],
+        ]
+    )
+    pg = PersistenceGradient(
+        homology_dimensions=[0, 1], zeta=0.0, collapse_edges=False, metric="precomputed"
+    )
+    assert all(
+        pg.phi(dist)
+        == torch.tensor(
+            [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 10.0, 10.0, 10.0, 10.0, 10.0]
+        )
+    )
+    assert pg.persistence_function(dist).item() > -23.0
     pg.sgd(dist, n_epochs=1, lr=0.002)
-                                  
+
 
 def test_persistence_gradient_pts():
     # test explicit gradients
-    pts = torch.tensor([[0., 0.],
-                        [0., 1.],
-                        [1., 0.]])
-    pg = PersistenceGradient(homology_dimensions=[0, 1],
-                             zeta=0.0,
-                             collapse_edges=False,
-                             metric="euclidean")
+    pts = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0]])
+    pg = PersistenceGradient(
+        homology_dimensions=[0, 1], zeta=0.0, collapse_edges=False, metric="euclidean"
+    )
 
-    assert all(pg.phi(pts)[:5] == torch.tensor([0., 0., 0.,
-                                                1., 1.]))
+    assert all(pg.phi(pts)[:5] == torch.tensor([0.0, 0.0, 0.0, 1.0, 1.0]))
 
     pg.sgd(pts, n_epochs=1, lr=0.002)
 
-    assert (pts.grad == torch.tensor([[1.,  1.],
-                                      [0., -1.],
-                                      [-1.,  0.]])).all().item()
+    assert (
+        (pts.grad == torch.tensor([[1.0, 1.0], [0.0, -1.0], [-1.0, 0.0]])).all().item()
+    )
diff --git a/gdeep/visualisation/tests/test_compactification.py b/gdeep/visualisation/tests/test_compactification.py
index 8d4b9f3d..44b1772b 100644
--- a/gdeep/visualisation/tests/test_compactification.py
+++ b/gdeep/visualisation/tests/test_compactification.py
@@ -1,4 +1,3 @@
-
 from gdeep.visualisation import Compactification
 from gdeep.models import FFNet
 from gdeep.utility import DEVICE
@@ -6,14 +5,15 @@
 
 def test_compactification():
     model = FFNet([2, 4, 2]).to(DEVICE)
-    c = Compactification(model,
-                         0.4,
-                         2,
-                         1500,
-                         0.05,
-                         50,
-                         [(-1., 1.), (-1., 1.)],
-                         )
+    c = Compactification(
+        model,
+        0.4,
+        2,
+        1500,
+        0.05,
+        50,
+        [(-1.0, 1.0), (-1.0, 1.0)],
+    )
     try:
         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
index 29768182..6becf62d 100644
--- a/gdeep/visualisation/tests/test_visualiser.py
+++ b/gdeep/visualisation/tests/test_visualiser.py
@@ -21,20 +21,19 @@
 transformed_textts = ptd.attach_transform_to_dataset(ds_val_str)  # type: ignore
 
 # 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)}
-                                                                  ])
+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)