diff --git a/.github/workflows/ci-build.yml b/.github/workflows/ci-build.yml
index 2ed4960d..3ffcb6f4 100644
--- a/.github/workflows/ci-build.yml
+++ b/.github/workflows/ci-build.yml
@@ -53,6 +53,7 @@ jobs:
           python -m pip install -r develop.txt
           python -m pip install -r docs/requirements.txt
           python -m pip install astropy scikit-image scikit-learn
+          python -m pip install tensorflow>=2.4.1
           python -m pip install twine
           python -m pip install .
 
diff --git a/MANIFEST.in b/MANIFEST.in
index 9a2f374e..74db0634 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -3,3 +3,4 @@ include develop.txt
 include docs/requirements.txt
 include README.rst
 include LICENSE.txt
+include docs/source/modopt_logo.png
diff --git a/README.md b/README.md
index aa72d976..0f7501f0 100644
--- a/README.md
+++ b/README.md
@@ -58,6 +58,7 @@ Note that none of these are required for running on a CPU.
 
 * [CuPy](https://cupy.dev/)
 * [Torch](https://pytorch.org/)
+* [TensorFlow](https://www.tensorflow.org/)
 
 ## Citation
 
diff --git a/docs/source/dependencies.rst b/docs/source/dependencies.rst
index 782807f8..2a513158 100644
--- a/docs/source/dependencies.rst
+++ b/docs/source/dependencies.rst
@@ -82,6 +82,7 @@ For GPU compliance the following packages can also be installed:
 
 * |link-to-cupy|
 * |link-to-torch|
+* |link-to-tf|
 
 .. |link-to-cupy| raw:: html
 
@@ -93,6 +94,11 @@ For GPU compliance the following packages can also be installed:
   <a href="https://pytorch.org/"
   target="_blank">Torch</a>
 
+.. |link-to-tf| raw:: html
+
+  <a href="https://www.tensorflow.org/"
+  target="_blank">TensorFlow</a>
+
 .. note::
 
   Note that none of these are required for running on a CPU.
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 8262d43d..238aa5b6 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -12,8 +12,8 @@ ModOpt Documentation
 .. include:: toc.rst
 
 :Author: Samuel Farrens `(samuel.farrens@cea.fr) <samuel.farrens@cea.fr>`_
-:Version: 1.5.0
-:Release Date: 31/03/2021
+:Version: 1.5.1
+:Release Date: 22/04/2021
 :Repository: |link-to-repo|
 
 .. |link-to-repo| raw:: html
diff --git a/modopt/base/backend.py b/modopt/base/backend.py
index 1f5a7a3a..5fbe912f 100644
--- a/modopt/base/backend.py
+++ b/modopt/base/backend.py
@@ -8,35 +8,73 @@
 
 """
 
-import warnings
 from importlib import util
 
 import numpy as np
 
+from modopt.interface.errors import warn
+
 try:
     import torch
+    from torch.utils.dlpack import from_dlpack as torch_from_dlpack
+    from torch.utils.dlpack import to_dlpack as torch_to_dlpack
+
 except ImportError:  # pragma: no cover
     import_torch = False
 else:
     import_torch = True
 
 # Handle the compatibility with variable
-gpu_compatibility = {
-    'cupy': False,
-    'cupy-cudnn': False,
+LIBRARIES = {
+    'cupy': None,
+    'tensorflow': None,
+    'numpy': np,
 }
 
 if util.find_spec('cupy') is not None:
     try:
         import cupy as cp
-        gpu_compatibility['cupy'] = True
+        LIBRARIES['cupy'] = cp
+    except ImportError:
+        pass
 
-        if util.find_spec('cupy.cuda.cudnn') is not None:
-            gpu_compatibility['cupy-cudnn'] = True
+if util.find_spec('tensorflow') is not None:
+    try:
+        from tensorflow.experimental import numpy as tnp
+        LIBRARIES['tensorflow'] = tnp
     except ImportError:
         pass
 
 
+def get_backend(backend):
+    """Get backend.
+
+    Returns the backend module for input specified by string
+
+    Parameters
+    ----------
+    backend: str
+        String holding the backend name. One of `tensorflow`,
+        `numpy` or `cupy`.
+
+    Returns
+    -------
+    tuple
+        Returns the module for carrying out calculations and the actual backend
+        that was reverted towards. If the right libraries are not installed,
+        the function warns and reverts to `numpy` backend
+    """
+    if backend not in LIBRARIES.keys() or LIBRARIES[backend] is None:
+        msg = (
+            '{0} backend not possible, please ensure that '
+            + 'the optional libraries are installed.\n'
+            + 'Reverting to numpy'
+        )
+        warn(msg.format(backend))
+        backend = 'numpy'
+    return LIBRARIES[backend], backend
+
+
 def get_array_module(input_data):
     """Get Array Module.
 
@@ -54,48 +92,47 @@ def get_array_module(input_data):
         The numpy or cupy module
 
     """
-    if gpu_compatibility['cupy']:
-        return cp.get_array_module(input_data)
-
+    if LIBRARIES['tensorflow'] is not None:
+        if isinstance(input_data, LIBRARIES['tensorflow'].ndarray):
+            return LIBRARIES['tensorflow']
+    if LIBRARIES['cupy'] is not None:
+        if isinstance(input_data, LIBRARIES['cupy'].ndarray):
+            return LIBRARIES['cupy']
     return np
 
 
-def move_to_device(input_data):
+def change_backend(input_data, backend='cupy'):
     """Move data to device.
 
-    This method moves data from CPU to GPU if we have the
-    compatibility to do so. It returns the same data if
-    it is already on GPU.
+    This method changes the backend of an array
+    This can be used to copy data to GPU or to CPU
 
     Parameters
     ----------
     input_data : numpy.ndarray or cupy.ndarray
         Input data array to be moved
+    backend: str, optional
+        The backend to use, one among `tensorflow`, `cupy` and
+        `numpy`. Default is `cupy`.
 
     Returns
     -------
-    cupy.ndarray
-        The CuPy array residing on GPU
+    backend.ndarray
+        An ndarray of specified backend
 
     """
     xp = get_array_module(input_data)
-
-    if xp == cp:
+    txp, target_backend = get_backend(backend)
+    if xp == txp:
         return input_data
-
-    if gpu_compatibility['cupy']:
-        return cp.array(input_data)
-
-    warnings.warn('Cupy is not installed, cannot move data to GPU')
-
-    return input_data
+    return txp.array(input_data)
 
 
 def move_to_cpu(input_data):
     """Move data to CPU.
 
-    This method moves data from GPU to CPU.It returns the same data if it is
-    already on CPU.
+    This method moves data from GPU to CPU.
+    It returns the same data if it is already on CPU.
 
     Parameters
     ----------
@@ -107,13 +144,20 @@ def move_to_cpu(input_data):
     numpy.ndarray
         The NumPy array residing on CPU
 
+    Raises
+    ------
+    ValueError
+        if the input does not correspond to any array
     """
     xp = get_array_module(input_data)
 
-    if xp == np:
+    if xp == LIBRARIES['numpy']:
         return input_data
-
-    return input_data.get()
+    elif xp == LIBRARIES['cupy']:
+        return input_data.get()
+    elif xp == LIBRARIES['tensorflow']:
+        return input_data.data.numpy()
+    raise ValueError('Cannot identify the array type.')
 
 
 def convert_to_tensor(input_data):
@@ -150,7 +194,7 @@ def convert_to_tensor(input_data):
     if xp == np:
         return torch.Tensor(input_data)
 
-    return torch.utils.dlpack.from_dlpack(input_data.toDlpack()).float()
+    return torch_from_dlpack(input_data.toDlpack()).float()
 
 
 def convert_to_cupy_array(input_data):
@@ -182,6 +226,6 @@ def convert_to_cupy_array(input_data):
         )
 
     if input_data.is_cuda:
-        return cp.fromDlpack(torch.utils.dlpack.to_dlpack(input_data))
+        return cp.fromDlpack(torch_to_dlpack(input_data))
 
     return input_data.detach().numpy()
diff --git a/modopt/math/matrix.py b/modopt/math/matrix.py
index cb54cebc..be737f52 100644
--- a/modopt/math/matrix.py
+++ b/modopt/math/matrix.py
@@ -12,12 +12,7 @@
 
 import numpy as np
 
-from modopt.base.backend import get_array_module
-
-try:
-    import cupy as cp
-except ImportError:  # pragma: no cover
-    pass
+from modopt.base.backend import get_array_module, get_backend
 
 
 def gram_schmidt(matrix, return_opt='orthonormal'):
@@ -303,7 +298,7 @@ def __init__(
         data_shape,
         data_type=float,
         auto_run=True,
-        use_gpu=False,
+        compute_backend='numpy',
         verbose=False,
     ):
 
@@ -311,10 +306,9 @@ def __init__(
         self._data_shape = data_shape
         self._data_type = data_type
         self._verbose = verbose
-        if use_gpu:
-            self.xp = cp
-        else:
-            self.xp = np
+        xp, compute_backend = get_backend(compute_backend)
+        self.xp = xp
+        self.compute_backend = compute_backend
         if auto_run:
             self.get_spec_rad()
 
diff --git a/modopt/opt/algorithms.py b/modopt/opt/algorithms.py
index ccd8fe21..125ac84c 100644
--- a/modopt/opt/algorithms.py
+++ b/modopt/opt/algorithms.py
@@ -54,11 +54,6 @@
 from modopt.opt.cost import costObj
 from modopt.opt.linear import Identity
 
-try:
-    import cupy as cp
-except ImportError:  # pragma: no cover
-    pass
-
 
 class SetUp(Observable):
     r"""Algorithm Set-Up.
@@ -92,7 +87,7 @@ def __init__(
         verbose=False,
         progress=True,
         step_size=None,
-        use_gpu=False,
+        compute_backend='numpy',
         **dummy_kwargs,
     ):
 
@@ -123,20 +118,9 @@ def __init__(
             )
             self.add_observer('cv_metrics', observer)
 
-        # Check for GPU
-        if use_gpu:
-            if backend.gpu_compatibility['cupy']:
-                self.xp = cp
-            else:
-                warn(
-                    'CuPy is not installed, cannot run on GPU!'
-                    + 'Running optimization on CPU.',
-                )
-                self.xp = np
-                use_gpu = False
-        else:
-            self.xp = np
-        self.use_gpu = use_gpu
+        xp, compute_backend = backend.get_backend(compute_backend)
+        self.xp = xp
+        self.compute_backend = compute_backend
 
     @property
     def metrics(self):
@@ -148,7 +132,9 @@ def metrics(self, metrics):
 
         if isinstance(metrics, type(None)):
             self._metrics = {}
-        elif not isinstance(metrics, dict):
+        elif isinstance(metrics, dict):
+            self._metrics = metrics
+        else:
             raise TypeError(
                 'Metrics must be a dictionary, not {0}.'.format(type(metrics)),
             )
@@ -184,10 +170,10 @@ def copy_data(self, input_data):
             Copy of input data
 
         """
-        if self.use_gpu:
-            return backend.move_to_device(input_data)
-
-        return self.xp.copy(input_data)
+        return self.xp.copy(backend.change_backend(
+            input_data,
+            self.compute_backend,
+        ))
 
     def _check_input_data(self, input_data):
         """Check input data type.
@@ -205,8 +191,10 @@ def _check_input_data(self, input_data):
             For invalid input type
 
         """
-        if not isinstance(input_data, self.xp.ndarray):
-            raise TypeError('Input data must be a numpy array.')
+        if not (isinstance(input_data, (self.xp.ndarray, np.ndarray))):
+            raise TypeError(
+                'Input data must be a numpy array or backend array',
+            )
 
     def _check_param(self, param_val):
         """Check algorithm parameters.
@@ -779,8 +767,8 @@ def _update(self):
             self._z_new = self._x_new
 
         # Update old values for next iteration.
-        self.xp.copyto(self._x_old, self._x_new)
-        self.xp.copyto(self._z_old, self._z_new)
+        self._x_old = self.xp.copy(self._x_new)
+        self._z_old = self.xp.copy(self._z_new)
 
         # Update parameter values for next iteration.
         self._update_param()
@@ -789,7 +777,7 @@ def _update(self):
         if self._cost_func:
             self.converge = (
                 self.any_convergence_flag()
-                or self._cost_func.get_cost(self._x_new),
+                or self._cost_func.get_cost(self._x_new)
             )
 
     def iterate(self, max_iter=150):
@@ -1548,7 +1536,7 @@ def _update(self):
         if self._cost_func:
             self.converge = (
                 self.any_convergence_flag()
-                or self._cost_func.get_cost(self._x_new),
+                or self._cost_func.get_cost(self._x_new)
             )
 
     def iterate(self, max_iter=150):
diff --git a/modopt/tests/test_base.py b/modopt/tests/test_base.py
index 26e1e4ea..873a4506 100644
--- a/modopt/tests/test_base.py
+++ b/modopt/tests/test_base.py
@@ -9,12 +9,14 @@
 """
 
 from builtins import range
-from unittest import TestCase
+from unittest import TestCase, skipIf
 
 import numpy as np
 import numpy.testing as npt
 
 from modopt.base import np_adjust, transform, types
+from modopt.base.backend import (LIBRARIES, change_backend, get_array_module,
+                                 get_backend)
 
 
 class NPAdjustTestCase(TestCase):
@@ -275,3 +277,53 @@ def test_check_npndarray(self):
             self.data3,
             dtype=np.integer,
         )
+
+
+class TestBackend(TestCase):
+    """Test the backend codes."""
+
+    def setUp(self):
+        """Set test parameter values."""
+        self.input = np.array([10, 10])
+
+    @skipIf(LIBRARIES['tensorflow'] is None, 'tensorflow library not installed')
+    def test_tf_backend(self):
+        """Test tensorflow backend."""
+        xp, backend = get_backend('tensorflow')
+        if backend != 'tensorflow' or xp != LIBRARIES['tensorflow']:
+            raise AssertionError('tensorflow get_backend fails!')
+        tf_input = change_backend(self.input, 'tensorflow')
+        if (
+            get_array_module(LIBRARIES['tensorflow'].ones(1)) != LIBRARIES['tensorflow']
+            or get_array_module(tf_input) != LIBRARIES['tensorflow']
+        ):
+            raise AssertionError('tensorflow backend fails!')
+
+    @skipIf(LIBRARIES['cupy'] is None, 'cupy library not installed')
+    def test_cp_backend(self):
+        """Test cupy backend."""
+        xp, backend = get_backend('cupy')
+        if backend != 'cupy' or xp != LIBRARIES['cupy']:
+            raise AssertionError('cupy get_backend fails!')
+        cp_input = change_backend(self.input, 'cupy')
+        if (
+            get_array_module(LIBRARIES['cupy'].ones(1)) != LIBRARIES['cupy']
+            or get_array_module(cp_input) != LIBRARIES['cupy']
+        ):
+            raise AssertionError('cupy backend fails!')
+
+    def test_np_backend(self):
+        """Test numpy backend."""
+        xp, backend = get_backend('numpy')
+        if backend != 'numpy' or xp != LIBRARIES['numpy']:
+            raise AssertionError('numpy get_backend fails!')
+        np_input = change_backend(self.input, 'numpy')
+        if (
+            get_array_module(LIBRARIES['numpy'].ones(1)) != LIBRARIES['numpy']
+            or get_array_module(np_input) != LIBRARIES['numpy']
+        ):
+            raise AssertionError('numpy backend fails!')
+
+    def tearDown(self):
+        """Tear Down of objects."""
+        self.input = None
diff --git a/modopt/tests/test_opt.py b/modopt/tests/test_opt.py
index b6805006..3c33c948 100644
--- a/modopt/tests/test_opt.py
+++ b/modopt/tests/test_opt.py
@@ -58,6 +58,17 @@ def setUp(self):
         reweight_inst = reweight.cwbReweight(self.data3)
         cost_inst = cost.costObj([grad_inst, prox_inst, prox_dual_inst])
         self.setup = algorithms.SetUp()
+        self.max_iter = 20
+
+        self.fb_all_iter = algorithms.ForwardBackward(
+            self.data1,
+            grad=grad_inst,
+            prox=prox_inst,
+            cost=None,
+            auto_iterate=False,
+            beta_update=func_identity,
+        )
+        self.fb_all_iter.iterate(self.max_iter)
 
         self.fb1 = algorithms.ForwardBackward(
             self.data1,
@@ -110,6 +121,17 @@ def setUp(self):
             s_greedy=1.1,
         )
 
+        self.gfb_all_iter = algorithms.GenForwardBackward(
+            self.data1,
+            grad=grad_inst,
+            prox_list=[prox_inst, prox_dual_inst],
+            cost=None,
+            auto_iterate=False,
+            gamma_update=func_identity,
+            beta_update=func_identity,
+        )
+        self.gfb_all_iter.iterate(self.max_iter)
+
         self.gfb1 = algorithms.GenForwardBackward(
             self.data1,
             grad=grad_inst,
@@ -133,6 +155,20 @@ def setUp(self):
             step_size=2,
         )
 
+        self.condat_all_iter = algorithms.Condat(
+            self.data1,
+            self.data2,
+            grad=grad_inst,
+            prox=prox_inst,
+            cost=None,
+            prox_dual=prox_dual_inst,
+            sigma_update=func_identity,
+            tau_update=func_identity,
+            rho_update=func_identity,
+            auto_iterate=False,
+        )
+        self.condat_all_iter.iterate(self.max_iter)
+
         self.condat1 = algorithms.Condat(
             self.data1,
             self.data2,
@@ -166,6 +202,18 @@ def setUp(self):
             auto_iterate=False,
         )
 
+        self.pogm_all_iter = algorithms.POGM(
+            u=self.data1,
+            x=self.data1,
+            y=self.data1,
+            z=self.data1,
+            grad=grad_inst,
+            prox=prox_inst,
+            auto_iterate=False,
+            cost=None,
+        )
+        self.pogm_all_iter.iterate(self.max_iter)
+
         self.pogm1 = algorithms.POGM(
             u=self.data1,
             x=self.data1,
@@ -184,13 +232,18 @@ def tearDown(self):
         self.data1 = None
         self.data2 = None
         self.setup = None
+        self.fb_all_iter = None
         self.fb1 = None
         self.fb2 = None
+        self.gfb_all_iter = None
         self.gfb1 = None
         self.gfb2 = None
+        self.condat_all_iter = None
         self.condat1 = None
         self.condat2 = None
         self.condat3 = None
+        self.pogm1 = None
+        self.pogm_all_iter = None
         self.dummy = None
 
     def test_set_up(self):
@@ -201,6 +254,17 @@ def test_set_up(self):
 
         npt.assert_raises(TypeError, self.setup._check_param_update, 1)
 
+    def test_all_iter(self):
+        """Test if all opt run for all iterations."""
+        opts = [
+            self.fb_all_iter,
+            self.gfb_all_iter,
+            self.condat_all_iter,
+            self.pogm_all_iter,
+        ]
+        for opt in opts:
+            npt.assert_equal(opt.idx, self.max_iter - 1)
+
     def test_forward_backward(self):
         """Test forward_backward."""
         npt.assert_array_equal(
diff --git a/setup.cfg b/setup.cfg
index 56a40d12..d2f544f0 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -26,7 +26,8 @@ per-file-ignores =
   #Justification: Needed for keeping package version and current API
   *__init__.py*: F401,F403,WPS347,WPS410,WPS412
   #Todo: Rethink conditional imports
-  modopt/base/backend.py: WPS229, WPS420
+  #Todo: How can we bypass mutable constants?
+  modopt/base/backend.py: WPS229, WPS420, WPS407
   #Todo: Rethink conditional imports
   modopt/base/observable.py: WPS420,WPS604
   #Todo: Check string for log formatting
@@ -54,6 +55,8 @@ per-file-ignores =
   modopt/signal/wavelet.py: S404,S603
   #Todo: Clean up tests
   modopt/tests/*.py: E731,F401,WPS301,WPS420,WPS425,WPS437,WPS604
+  #Todo: Import has bad parenthesis
+  modopt/tests/test_base.py: WPS318,WPS319,E501,WPS301
 #WPS Settings
 max-arguments = 25
 max-attributes = 40
diff --git a/setup.py b/setup.py
index 92040f6d..841ca1b1 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@
 # Set the package release version
 major = 1
 minor = 5
-patch = 0
+patch = 1
 
 # Set the package details
 name = 'modopt'
@@ -29,9 +29,10 @@
 os_str = 'Operating System :: {0}'
 
 classifiers = (
-    [lc_str.format(license)] + [ln_str] +
-    [py_str.format(ver) for ver in python_versions_supported] +
-    [os_str.format(ops) for ops in os_platforms_supported]
+    [lc_str.format(license)]
+    + [ln_str]
+    + [py_str.format(ver) for ver in python_versions_supported]
+    + [os_str.format(ops) for ops in os_platforms_supported]
 )
 
 # Source package description from README.md