diff --git a/src/qiboml/backends/tensorflow.py b/src/qiboml/backends/tensorflow.py
index 9e1a3bd..7d66a60 100644
--- a/src/qiboml/backends/tensorflow.py
+++ b/src/qiboml/backends/tensorflow.py
@@ -2,7 +2,6 @@
 import os
 
 import numpy as np
-
 from qibo import __version__
 from qibo.backends.npmatrices import NumpyMatrices
 from qibo.backends.numpy import NumpyBackend
@@ -14,7 +13,6 @@ class TensorflowMatrices(NumpyMatrices):
 
     def __init__(self, dtype):
         super().__init__(dtype)
-
         import tensorflow as tf  # pylint: disable=import-error
         import tensorflow.experimental.numpy as tnp  # pylint: disable=import-error
 
@@ -95,9 +93,6 @@ def to_numpy(self, x):
     def compile(self, func):
         return self.tf.function(func)
 
-    def customize_gradient(self, func):
-        return self.tf.custom_gradient(func)
-
     def zero_state(self, nqubits):
         idx = self.tf.constant([[0]], dtype="int32")
         state = self.tf.zeros((2**nqubits,), dtype=self.dtype)
@@ -203,18 +198,16 @@ def calculate_hamiltonian_state_product(self, matrix, state):
         else:
             raise_error(
                 ValueError,
-                "Cannot multiply Hamiltonian with " "rank-{} tensor.".format(rank),
+                f"Cannot multiply Hamiltonian with rank-{rank} tensor.",
             )
 
-        return expectation_with_custom_gradient(params)
-
     def test_regressions(self, name):
         if name == "test_measurementresult_apply_bitflips":
             return [
-                [4, 0, 0, 1, 0, 2, 2, 4, 4, 0],
-                [4, 0, 0, 1, 0, 2, 2, 4, 4, 0],
-                [4, 0, 0, 1, 0, 0, 0, 4, 4, 0],
-                [4, 0, 0, 0, 0, 0, 0, 4, 4, 0],
+                [4, 0, 0, 1, 0, 0, 1, 0, 0, 0],
+                [0, 1, 1, 2, 1, 1, 4, 0, 0, 4],
+                [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 4, 0, 0, 0, 4],
             ]
         elif name == "test_probabilistic_measurement":
             if "GPU" in self.device:  # pragma: no cover
@@ -229,6 +222,6 @@ def test_regressions(self, name):
         elif name == "test_post_measurement_bitflips_on_circuit":
             return [
                 {5: 30},
-                {5: 16, 7: 10, 6: 2, 3: 1, 4: 1},
-                {3: 6, 5: 6, 7: 5, 2: 4, 4: 3, 0: 2, 1: 2, 6: 2},
+                {5: 12, 7: 6, 4: 6, 1: 5, 6: 1},
+                {3: 7, 6: 4, 2: 4, 7: 4, 0: 4, 5: 3, 4: 2, 1: 2},
             ]