diff --git a/src/qibo/core/callbacks.py b/src/qibo/core/callbacks.py
index 22cb22d58b..af2e0838bb 100644
--- a/src/qibo/core/callbacks.py
+++ b/src/qibo/core/callbacks.py
@@ -50,6 +50,9 @@ def density_matrix(self, x):
     @abstract_callbacks.Callback.nqubits.setter
     def nqubits(self, n: int):
         from qibo import gates
+        if self._nqubits is not None and self._nqubits != n:
+            raise_error(RuntimeError,
+                        f"Changing EntanglementEntropy nqubits from {self._nqubits} to {n}.")
         self._nqubits = n
         if self.partition is None:
             self.partition = list(range(n // 2 + n % 2))
@@ -77,9 +80,7 @@ def set_nqubits(self, state):
         if not isinstance(state, K.tensor_types):
             raise_error(TypeError, "State of unknown type {} was given in callback "
                                    "calculation.".format(type(state)))
-        if self._nqubits is None:
-            self.nqubits = int(math.log2(tuple(state.shape)[0]))
-
+        self.nqubits = int(math.log2(tuple(state.shape)[0]))
 
     def state_vector_call(self, state):
         self.set_nqubits(state)
diff --git a/src/qibo/tests/test_core_callbacks.py b/src/qibo/tests/test_core_callbacks.py
index c9eda08e37..79264323e8 100644
--- a/src/qibo/tests/test_core_callbacks.py
+++ b/src/qibo/tests/test_core_callbacks.py
@@ -195,6 +195,11 @@ def target_entropy(t):
     target = [0, target_entropy(0.1234), 0, target_entropy(0.4321)]
     K.assert_allclose(entropy[:], target, atol=_atol)
 
+    c = Circuit(8, accelerators)
+    with pytest.raises(RuntimeError):
+        c.add(gates.CallbackGate(entropy))
+        state = c()
+
 
 def test_entropy_large_circuit(backend, accelerators):
     """Check that entropy calculation works for variational like circuit."""